[WebAssembly] Improve libObject support for wasm imports and exports

Previously we had only supported the importing and
exporting of functions and globals.

Also, add usefull overload of getWasmSymbol() and
getNumberOfSymbols() in support of lld port.

Differential Revision: https://reviews.llvm.org/D33011

llvm-svn: 302601
diff --git a/llvm/include/llvm/Object/Wasm.h b/llvm/include/llvm/Object/Wasm.h
index 4bc39d9..d200d4a 100644
--- a/llvm/include/llvm/Object/Wasm.h
+++ b/llvm/include/llvm/Object/Wasm.h
@@ -67,7 +67,8 @@
   WasmObjectFile(MemoryBufferRef Object, Error &Err);
 
   const wasm::WasmObjectHeader &getHeader() const;
-  const WasmSymbol &getWasmSymbol(DataRefImpl Symb) const;
+  const WasmSymbol &getWasmSymbol(const DataRefImpl &Symb) const;
+  const WasmSymbol &getWasmSymbol(const SymbolRef &Symbol) const;
   const WasmSection &getWasmSection(const SectionRef &Section) const;
   const wasm::WasmRelocation &getWasmRelocation(const RelocationRef& Ref) const;
 
@@ -81,6 +82,10 @@
   const std::vector<wasm::WasmGlobal>& globals() const { return Globals; }
   const std::vector<wasm::WasmExport>& exports() const { return Exports; }
 
+  uint32_t getNumberOfSymbols() const {
+    return Symbols.size();
+  }
+
   const std::vector<wasm::WasmElemSegment>& elements() const {
     return ElemSegments;
   }
diff --git a/llvm/include/llvm/ObjectYAML/WasmYAML.h b/llvm/include/llvm/ObjectYAML/WasmYAML.h
index bd7d72b..747cd62 100644
--- a/llvm/include/llvm/ObjectYAML/WasmYAML.h
+++ b/llvm/include/llvm/ObjectYAML/WasmYAML.h
@@ -34,17 +34,6 @@
   yaml::Hex32 Version;
 };
 
-struct Import {
-  StringRef Module;
-  StringRef Field;
-  ExportKind Kind;
-  union {
-    uint32_t SigIndex;
-    ValueType GlobalType;
-  };
-  bool GlobalMutable;
-};
-
 struct Limits {
   yaml::Hex32 Flags;
   yaml::Hex32 Initial;
@@ -74,6 +63,18 @@
   wasm::WasmInitExpr InitExpr;
 };
 
+struct Import {
+  StringRef Module;
+  StringRef Field;
+  ExportKind Kind;
+  union {
+    uint32_t SigIndex;
+    Global Global;
+    Table Table;
+    Limits Memory;
+  };
+};
+
 struct LocalDecl {
   ValueType Type;
   uint32_t Count;
diff --git a/llvm/include/llvm/Support/Wasm.h b/llvm/include/llvm/Support/Wasm.h
index a48dfe10..e383182 100644
--- a/llvm/include/llvm/Support/Wasm.h
+++ b/llvm/include/llvm/Support/Wasm.h
@@ -37,17 +37,6 @@
   int32_t ReturnType;
 };
 
-struct WasmImport {
-  StringRef Module;
-  StringRef Field;
-  uint32_t Kind;
-  union {
-    uint32_t SigIndex;
-    int32_t GlobalType;
-  };
-  bool GlobalMutable;
-};
-
 struct WasmExport {
   StringRef Name;
   uint32_t Kind;
@@ -82,6 +71,18 @@
   WasmInitExpr InitExpr;
 };
 
+struct WasmImport {
+  StringRef Module;
+  StringRef Field;
+  uint32_t Kind;
+  union {
+    uint32_t SigIndex;
+    WasmGlobal Global;
+    WasmTable Table;
+    WasmLimits Memory;
+  };
+};
+
 struct WasmLocalDecl {
   int32_t Type;
   uint32_t Count;
diff --git a/llvm/lib/Object/WasmObjectFile.cpp b/llvm/lib/Object/WasmObjectFile.cpp
index 012c9dc..058686e 100644
--- a/llvm/lib/Object/WasmObjectFile.cpp
+++ b/llvm/lib/Object/WasmObjectFile.cpp
@@ -168,6 +168,13 @@
   return Result;
 }
 
+static wasm::WasmTable readTable(const uint8_t *&Ptr) {
+  wasm::WasmTable Table;
+  Table.ElemType = readVarint7(Ptr);
+  Table.Limits = readLimits(Ptr);
+  return Table;
+}
+
 static Error readSection(WasmSection &Section, const uint8_t *&Ptr,
                          const uint8_t *Start) {
   // TODO(sbc): Avoid reading past EOF in the case of malformed files.
@@ -397,13 +404,22 @@
                            Sections.size(), i);
       break;
     case wasm::WASM_EXTERNAL_GLOBAL:
-      Im.GlobalType = readVarint7(Ptr);
-      Im.GlobalMutable = readVaruint1(Ptr);
+      Im.Global.Type = readVarint7(Ptr);
+      Im.Global.Mutable = readVaruint1(Ptr);
       Symbols.emplace_back(Im.Field, WasmSymbol::SymbolType::GLOBAL_IMPORT,
                            Sections.size(), i);
       break;
+    case wasm::WASM_EXTERNAL_MEMORY:
+      Im.Memory = readLimits(Ptr);
+      break;
+    case wasm::WASM_EXTERNAL_TABLE:
+      Im.Table = readTable(Ptr);
+      if (Im.Table.ElemType != wasm::WASM_TYPE_ANYFUNC) {
+        return make_error<GenericBinaryError>("Invalid table element type",
+                                              object_error::parse_failed);
+      }
+      break;
     default:
-      // TODO(sbc): Handle other kinds of imports
       return make_error<GenericBinaryError>(
           "Unexpected import kind", object_error::parse_failed);
     }
@@ -431,14 +447,11 @@
   uint32_t Count = readVaruint32(Ptr);
   Tables.reserve(Count);
   while (Count--) {
-    wasm::WasmTable Table;
-    Table.ElemType = readVarint7(Ptr);
-    if (Table.ElemType != wasm::WASM_TYPE_ANYFUNC) {
+    Tables.push_back(readTable(Ptr));
+    if (Tables.back().ElemType != wasm::WASM_TYPE_ANYFUNC) {
       return make_error<GenericBinaryError>("Invalid table element type",
                                             object_error::parse_failed);
     }
-    Table.Limits = readLimits(Ptr);
-    Tables.push_back(Table);
   }
   if (Ptr != End)
     return make_error<GenericBinaryError>("Table section ended prematurely",
@@ -493,8 +506,10 @@
       Symbols.emplace_back(Ex.Name, WasmSymbol::SymbolType::GLOBAL_EXPORT,
                            Sections.size(), i);
       break;
+    case wasm::WASM_EXTERNAL_MEMORY:
+    case wasm::WASM_EXTERNAL_TABLE:
+      break;
     default:
-      // TODO(sbc): Handle other kinds of exports
       return make_error<GenericBinaryError>(
           "Unexpected export kind", object_error::parse_failed);
     }
@@ -638,10 +653,14 @@
   return BasicSymbolRef(Ref, this);
 }
 
-const WasmSymbol &WasmObjectFile::getWasmSymbol(DataRefImpl Symb) const {
+const WasmSymbol &WasmObjectFile::getWasmSymbol(const DataRefImpl &Symb) const {
   return Symbols[Symb.d.a];
 }
 
+const WasmSymbol &WasmObjectFile::getWasmSymbol(const SymbolRef &Symb) const {
+  return getWasmSymbol(Symb.getRawDataRefImpl());
+}
+
 Expected<StringRef> WasmObjectFile::getSymbolName(DataRefImpl Symb) const {
   const WasmSymbol &Sym = getWasmSymbol(Symb);
   return Sym.Name;
diff --git a/llvm/lib/ObjectYAML/WasmYAML.cpp b/llvm/lib/ObjectYAML/WasmYAML.cpp
index c5d1b43..514ae55 100644
--- a/llvm/lib/ObjectYAML/WasmYAML.cpp
+++ b/llvm/lib/ObjectYAML/WasmYAML.cpp
@@ -265,8 +265,12 @@
   if (Import.Kind == wasm::WASM_EXTERNAL_FUNCTION) {
     IO.mapRequired("SigIndex", Import.SigIndex);
   } else if (Import.Kind == wasm::WASM_EXTERNAL_GLOBAL) {
-    IO.mapRequired("GlobalType", Import.GlobalType);
-    IO.mapRequired("GlobalMutable", Import.GlobalMutable);
+    IO.mapRequired("GlobalType", Import.Global.Type);
+    IO.mapRequired("GlobalMutable", Import.Global.Mutable);
+  } else if (Import.Kind == wasm::WASM_EXTERNAL_TABLE) {
+    IO.mapRequired("Table", Import.Table);
+  } else if (Import.Kind == wasm::WASM_EXTERNAL_MEMORY ) {
+    IO.mapRequired("Memory", Import.Memory);
   } else {
     llvm_unreachable("unhandled import type");
   }
diff --git a/llvm/test/ObjectYAML/wasm/export_section.yaml b/llvm/test/ObjectYAML/wasm/export_section.yaml
index 1d1a16f..89ebee3 100644
--- a/llvm/test/ObjectYAML/wasm/export_section.yaml
+++ b/llvm/test/ObjectYAML/wasm/export_section.yaml
@@ -5,12 +5,18 @@
 Sections:
   - Type:            EXPORT
     Exports:         
-      - Name:            foo
-        Kind:            FUNCTION
-        Index:           0
-      - Name:            bar
+      - Name:            function_export
         Kind:            FUNCTION
         Index:           1
+      - Name:            global_export
+        Kind:            GLOBAL
+        Index:           1
+      - Name:            memory_export
+        Kind:            MEMORY
+        Index:           0
+      - Name:            table_export
+        Kind:            TABLE
+        Index:           0
 ...
 # CHECK: --- !WASM
 # CHECK: FileHeader:
@@ -18,10 +24,16 @@
 # CHECK: Sections:
 # CHECK:   - Type:            EXPORT
 # CHECK:     Exports:         
-# CHECK:       - Name:            foo
-# CHECK:         Kind:            FUNCTION
-# CHECK:         Index:           0
-# CHECK:       - Name:            bar
+# CHECK:       - Name:            function_export
 # CHECK:         Kind:            FUNCTION
 # CHECK:         Index:           1
+# CHECK:       - Name:            global_export
+# CHECK:         Kind:            GLOBAL
+# CHECK:         Index:           1
+# CHECK:       - Name:            memory_export
+# CHECK:         Kind:            MEMORY
+# CHECK:         Index:           0
+# CHECK:       - Name:            table_export
+# CHECK:         Kind:            TABLE
+# CHECK:         Index:           0
 # CHECK: ...
diff --git a/llvm/test/ObjectYAML/wasm/import_section.yaml b/llvm/test/ObjectYAML/wasm/import_section.yaml
index 52f466a..115d4cc 100644
--- a/llvm/test/ObjectYAML/wasm/import_section.yaml
+++ b/llvm/test/ObjectYAML/wasm/import_section.yaml
@@ -9,19 +9,32 @@
         ParamTypes:
           - I32
   - Type:            IMPORT
-    Imports:         
+    Imports:
       - Module:          foo
-        Field:           bar
+        Field:           imported_function
         Kind:            FUNCTION
         SigIndex:        0
       - Module:          fiz
-        Field:           baz
+        Field:           imported_global
         Kind:            GLOBAL
         GlobalType:      I32
         GlobalMutable:   false
-  - Type:            FUNCTION
-    FunctionTypes:
-      - 0
+      - Module:          foo
+        Field:           imported_memory
+        Kind:            MEMORY
+        Memory:
+          Flags:           0x00000001
+          Initial:         0x00000010
+          Maximum:         0x00000011
+      - Module:          foo
+        Field:           imported_table
+        Kind:            TABLE
+        Table:
+          ElemType:      ANYFUNC
+          Limits:
+            Flags:           0x00000001
+            Initial:         0x00000020
+            Maximum:         0x00000022
 ...
 # CHECK: --- !WASM
 # CHECK: FileHeader:
@@ -30,12 +43,28 @@
 # CHECK:   - Type:            IMPORT
 # CHECK:     Imports:         
 # CHECK:       - Module:          foo
-# CHECK:         Field:           bar
+# CHECK:         Field:           imported_function
 # CHECK:         Kind:            FUNCTION
 # CHECK:         SigIndex:        0
 # CHECK:       - Module:          fiz
-# CHECK:         Field:           baz
+# CHECK:         Field:           imported_global
 # CHECK:         Kind:            GLOBAL
 # CHECK:         GlobalType:      I32
 # CHECK:         GlobalMutable:   false
+# CHECK:       - Module:          foo
+# CHECK:         Field:           imported_memory
+# CHECK:         Kind:            MEMORY
+# CHECK:         Memory:
+# CHECK:           Flags:           0x00000001
+# CHECK:           Initial:         0x00000010
+# CHECK:           Maximum:         0x00000011
+# CHECK:       - Module:          foo
+# CHECK:         Field:           imported_table
+# CHECK:         Kind:            TABLE
+# CHECK:         Table:
+# CHECK:           ElemType:      ANYFUNC
+# CHECK:           Limits:
+# CHECK:             Flags:           0x00000001
+# CHECK:             Initial:         0x00000020
+# CHECK:             Maximum:         0x00000022
 # CHECK: ...
diff --git a/llvm/tools/obj2yaml/wasm2yaml.cpp b/llvm/tools/obj2yaml/wasm2yaml.cpp
index cc04b99..6efa2ac 100644
--- a/llvm/tools/obj2yaml/wasm2yaml.cpp
+++ b/llvm/tools/obj2yaml/wasm2yaml.cpp
@@ -25,6 +25,23 @@
   ErrorOr<WasmYAML::Object *> dump();
 };
 
+WasmYAML::Table make_table(const wasm::WasmTable &Table) {
+  WasmYAML::Table T;
+  T.ElemType = Table.ElemType;
+  T.TableLimits.Flags = Table.Limits.Flags;
+  T.TableLimits.Initial = Table.Limits.Initial;
+  T.TableLimits.Maximum = Table.Limits.Maximum;
+  return T;
+}
+
+WasmYAML::Limits make_limits(const wasm::WasmLimits &Limits) {
+  WasmYAML::Limits L;
+  L.Flags = Limits.Flags;
+  L.Initial = Limits.Initial;
+  L.Maximum = Limits.Maximum;
+  return L;
+}
+
 ErrorOr<WasmYAML::Object *> WasmDumper::dump() {
   auto Y = make_unique<WasmYAML::Object>();
 
@@ -82,17 +99,26 @@
     case wasm::WASM_SEC_IMPORT: {
       auto ImportSec = make_unique<WasmYAML::ImportSection>();
       for (auto &Import : Obj.imports()) {
-        WasmYAML::Import Ex;
-        Ex.Module = Import.Module;
-        Ex.Field = Import.Field;
-        Ex.Kind = Import.Kind;
-        if (Ex.Kind == wasm::WASM_EXTERNAL_FUNCTION) {
-          Ex.SigIndex = Import.SigIndex;
-        } else if (Ex.Kind == wasm::WASM_EXTERNAL_GLOBAL) {
-          Ex.GlobalType = Import.GlobalType;
-          Ex.GlobalMutable = Import.GlobalMutable;
+        WasmYAML::Import Im;
+        Im.Module = Import.Module;
+        Im.Field = Import.Field;
+        Im.Kind = Import.Kind;
+        switch (Im.Kind) {
+        case wasm::WASM_EXTERNAL_FUNCTION:
+          Im.SigIndex = Import.SigIndex;
+          break;
+        case wasm::WASM_EXTERNAL_GLOBAL:
+          Im.Global.Type = Import.Global.Type;
+          Im.Global.Mutable = Import.Global.Mutable;
+          break;
+        case wasm::WASM_EXTERNAL_TABLE:
+          Im.Table = make_table(Import.Table);
+          break;
+        case wasm::WASM_EXTERNAL_MEMORY:
+          Im.Memory = make_limits(Import.Memory);
+          break;
         }
-        ImportSec->Imports.push_back(Ex);
+        ImportSec->Imports.push_back(Im);
       }
       S = std::move(ImportSec);
       break;
@@ -107,25 +133,16 @@
     }
     case wasm::WASM_SEC_TABLE: {
       auto TableSec = make_unique<WasmYAML::TableSection>();
-      for (auto &Table : Obj.tables()) {
-        WasmYAML::Table T;
-        T.ElemType = Table.ElemType;
-        T.TableLimits.Flags = Table.Limits.Flags;
-        T.TableLimits.Initial = Table.Limits.Initial;
-        T.TableLimits.Maximum = Table.Limits.Maximum;
-        TableSec->Tables.push_back(T);
+      for (const wasm::WasmTable &Table : Obj.tables()) {
+        TableSec->Tables.push_back(make_table(Table));
       }
       S = std::move(TableSec);
       break;
     }
     case wasm::WASM_SEC_MEMORY: {
       auto MemorySec = make_unique<WasmYAML::MemorySection>();
-      for (auto &Memory : Obj.memories()) {
-        WasmYAML::Limits L;
-        L.Flags = Memory.Flags;
-        L.Initial = Memory.Initial;
-        L.Maximum = Memory.Maximum;
-        MemorySec->Memories.push_back(L);
+      for (const wasm::WasmLimits &Memory : Obj.memories()) {
+        MemorySec->Memories.push_back(make_limits(Memory));
       }
       S = std::move(MemorySec);
       break;
diff --git a/llvm/tools/yaml2obj/yaml2wasm.cpp b/llvm/tools/yaml2obj/yaml2wasm.cpp
index eed9f2c..502b1e6 100644
--- a/llvm/tools/yaml2obj/yaml2wasm.cpp
+++ b/llvm/tools/yaml2obj/yaml2wasm.cpp
@@ -169,8 +169,15 @@
       encodeULEB128(Import.SigIndex, OS);
       break;
     case wasm::WASM_EXTERNAL_GLOBAL:
-      encodeSLEB128(Import.GlobalType, OS);
-      writeUint8(OS, Import.GlobalMutable);
+      encodeSLEB128(Import.Global.Type, OS);
+      writeUint8(OS, Import.Global.Mutable);
+      break;
+    case wasm::WASM_EXTERNAL_MEMORY:
+      writeLimits(Import.Memory, OS);
+      break;
+    case wasm::WASM_EXTERNAL_TABLE:
+      encodeSLEB128(Import.Table.ElemType, OS);
+      writeLimits(Import.Table.TableLimits, OS);
       break;
     default:
       errs() << "Unknown import type: " << Import.Kind;