Fix import section count when using compact imports

include/wabt/binary-reader-objdump.h

@@ -87,6 +87,9 @@ struct ObjdumpState {
   std::vector<ObjdumpSymbol> symtab;
   std::map<Index, Index> function_param_counts;
   std::map<Index, Index> function_types;
+  // Total number of imports.  This can differ from the import section count
+  // when compact imports are used.
+  Index num_imports = 0;
 };
 
 Result ReadBinaryObjdump(ByteSpan data,

include/wabt/binary-reader.h

@@ -115,6 +115,9 @@ class BinaryReaderDelegate {
 
   /* Import section */
   virtual Result BeginImportSection(Offset size) = 0;
+  // `count` is the number of entries in the import section.  When compact
+  // imports are used this can be less than the number of imports, since a
+  // single entry can contain multiple imports.
   virtual Result OnImportCount(Index count) = 0;
   virtual Result OnImport(Index index,
                           ExternalKind kind,

src/binary-reader-objdump.cc

@@ -379,6 +379,14 @@ class BinaryReaderObjdumpPrepass : public BinaryReaderObjdumpBase {
     return Result::Ok;
   }
 
+  Result OnImport(Index index,
+                  ExternalKind kind,
+                  std::string_view module_name,
+                  std::string_view field_name) override {
+    objdump_state_->num_imports = index + 1;
+    return Result::Ok;
+  }
+
   Result OnImportFunc(Index import_index,
                       std::string_view module_name,
                       std::string_view field_name,
@@ -1580,7 +1588,10 @@ Result BinaryReaderObjdump::OnDataCount(Index count) {
 }
 
 Result BinaryReaderObjdump::OnImportCount(Index count) {
-  return OnCount(count);
+  // `count` is the number of entries in the import section, which can be
+  // less than the number of imports when compact imports are used.  Report
+  // the total number of imports (as counted during the prepass) instead.
+  return OnCount(std::max(count, objdump_state_->num_imports));
 }
 
 Result BinaryReaderObjdump::OnImportFunc(Index import_index,

src/binary-reader.cc

@@ -169,8 +169,7 @@ class BinaryReader {
   Result ReadCodeMetadataSection(std::string_view name, Offset section_size);
   Result ReadCustomSection(Index section_index, Offset section_size);
   Result ReadTypeSection(Offset section_size);
-  Result ReadImport(Index i,
-                    std::string_view module_name,
+  Result ReadImport(std::string_view module_name,
                     std::string_view field_name,
                     ExternalKind kind);
   Result ReadImportSection(Offset section_size);
@@ -2737,10 +2736,13 @@ Result BinaryReader::ReadTypeSection(Offset section_size) {
   return Result::Ok;
 }
 
-Result BinaryReader::ReadImport(Index i,
-                                std::string_view module_name,
+Result BinaryReader::ReadImport(std::string_view module_name,
                                 std::string_view field_name,
                                 ExternalKind kind) {
+  // Each import increments exactly one of the per-kind counters below, so
+  // their sum is the index of this import within the import section.
+  Index i = num_func_imports_ + num_table_imports_ + num_memory_imports_ +
+            num_global_imports_ + num_tag_imports_;
   CALLBACK(OnImport, i, kind, module_name, field_name);
   switch (kind) {
     case ExternalKind::Func: {
@@ -2800,11 +2802,13 @@ Result BinaryReader::ReadImport(Index i,
 
 Result BinaryReader::ReadImportSection(Offset section_size) {
   CALLBACK(BeginImportSection, section_size);
-  Index num_imports;
-  CHECK_RESULT(ReadCount(&num_imports, "import count"));
-  CALLBACK(OnImportCount, num_imports);
-  Index i = 0;
-  while (i < num_imports) {
+  // Note: With compact imports a single entry in the import section can
+  // contain more than one import, so the number of entries is not necessarily
+  // the same as the number of imports.
+  Index num_entries;
+  CHECK_RESULT(ReadCount(&num_entries, "import count"));
+  CALLBACK(OnImportCount, num_entries);
+  for (Index entry = 0; entry < num_entries; ++entry) {
     std::string_view module_name;
     CHECK_RESULT(ReadStr(&module_name, "import module name"));
     std::string_view field_name;
@@ -2824,15 +2828,15 @@ Result BinaryReader::ReadImportSection(Offset section_size) {
         CHECK_RESULT(ReadCount(&num_compact_imports, "compact import count"));
         for (Index j = 0; j < num_compact_imports; ++j) {
           CHECK_RESULT(ReadStr(&field_name, "compact import field name"));
-          CHECK_RESULT(ReadImport(i++, module_name, field_name, kind));
+          CHECK_RESULT(ReadImport(module_name, field_name, kind));
         }
       } else {
         CHECK_RESULT(ReadCount(&num_compact_imports, "compact import count"));
         for (Index j = 0; j < num_compact_imports; ++j) {
           CHECK_RESULT(ReadStr(&field_name, "compact import field name"));
           CHECK_RESULT(
               ReadExternalKind(&kind, "compact import kind", "import"));
-          CHECK_RESULT(ReadImport(i++, module_name, field_name, kind));
+          CHECK_RESULT(ReadImport(module_name, field_name, kind));
         }
       }
     } else {
@@ -2841,7 +2845,7 @@ Result BinaryReader::ReadImportSection(Offset section_size) {
       // byte so we can read it with ReadExternalKind
       state_.offset--;
       CHECK_RESULT(ReadExternalKind(&kind, "import kind", "import"));
-      CHECK_RESULT(ReadImport(i++, module_name, field_name, kind));
+      CHECK_RESULT(ReadImport(module_name, field_name, kind));
     }
   }
 

src/binary-writer.cc

@@ -1497,45 +1497,56 @@ Result BinaryWriter::WriteModule() {
 
   if (module_->imports.size()) {
     BeginKnownSection(BinarySection::Import);
-    WriteU32Leb128(stream_, module_->imports.size(), "num imports");
+
+    // Returns the number of consecutive imports, starting at index `i`, that
+    // share the same module name.  When compact imports are enabled, a run of
+    // more than one such import is written as a single (compact) entry.
+    auto get_group_size = [&](size_t i) -> size_t {
+      size_t group_size = 1;
+      if (options_.features.compact_imports_enabled()) {
+        // Currently we only support grouping by module name (0x7F mode)
+        // and not the module name + kind grouping (0x7E mode).
+        const std::string& module_name = module_->imports[i]->module_name;
+        while (i + group_size < module_->imports.size() &&
+               module_->imports[i + group_size]->module_name == module_name) {
+          group_size++;
+        }
+      }
+      return group_size;
+    };
+
+    // The import section count is the number of entries in the section, not
+    // the number of imports, since a single compact entry can contain many
+    // imports.
+    size_t num_entries = 0;
+    for (size_t i = 0; i < module_->imports.size(); i += get_group_size(i)) {
+      num_entries++;
+    }
+    WriteU32Leb128(stream_, num_entries, "num imports");
 
     size_t i = 0;
     while (i < module_->imports.size()) {
       const Import* import = module_->imports[i];
       WriteHeader("import header", i);
       WriteStr(stream_, import->module_name, "import module name",
                PrintChars::Yes);
-      bool compact = false;
-      if (options_.features.compact_imports_enabled()) {
-        // Write compact imports when they are available.
-        // Currently we only support grouping by module name (0x7F mode)
-        // and not the module name + kind grouping (0x7E mode).
-        size_t group_size = 1;
-        size_t j = i + 1;
-        while (j < module_->imports.size() &&
-               import->module_name == module_->imports[j]->module_name) {
-          group_size++;
-          j++;
+      // Use compact imports iff we have a continuous sequence of more than
+      // one import with the same module name.
+      size_t group_size = get_group_size(i);
+      if (group_size > 1) {
+        WriteStr(stream_, "", "empty field name", PrintChars::Yes);
+        stream_->WriteU8(0x7F, "compact import marker");
+        WriteU32Leb128(stream_, group_size, "import group size");
+        while (group_size--) {
+          WriteHeader("compact import header", i);
+          const Import* import = module_->imports[i];
+          WriteStr(stream_, import->field_name, "import field name",
+                   PrintChars::Yes);
+          stream_->WriteU8Enum(import->kind(), "import kind");
+          WriteImport(import);
+          i++;
         }
-        // Use compact imports iff we have a continuous sequence of more than
-        // one import with the same module name.
-        if (group_size > 1) {
-          compact = true;
-          WriteStr(stream_, "", "empty field name", PrintChars::Yes);
-          stream_->WriteU8(0x7F, "compact import marker");
-          WriteU32Leb128(stream_, group_size, "import group size");
-          while (group_size--) {
-            WriteHeader("compact import header", i);
-            const Import* import = module_->imports[i];
-            WriteStr(stream_, import->field_name, "import field name",
-                     PrintChars::Yes);
-            stream_->WriteU8Enum(import->kind(), "import kind");
-            WriteImport(import);
-            i++;
-          }
-        }
-      }
-      if (!compact) {
+      } else {
         WriteStr(stream_, import->field_name, "import field name",
                  PrintChars::Yes);
         stream_->WriteU8Enum(import->kind(), "import kind");

test/binary/compact-imports.txt

@@ -6,7 +6,7 @@ magic
 version
 section(TYPE) { count[1] function params[0] results[1] i32 }
 section(IMPORT) {
-  count[6]
+  count[4]
   str("mod") str("func") func_kind type[0]
   str("mod") str("") compact[0x7F]
   count[3]
@@ -18,15 +18,16 @@ section(IMPORT) {
   count[2]
   str("func5") type[0]
   str("func6") type[0]
+  str("mod2") str("func7") func_kind type[0]
 }
 (;; STDERR ;;;
 BeginModule(version: 1)
   BeginTypeSection(5)
     OnTypeCount(1)
     OnFuncType(index: 0, params: [], results: [i32])
   EndTypeSection
-  BeginImportSection(65)
-    OnImportCount(6)
+  BeginImportSection(78)
+    OnImportCount(4)
     OnImport(index: 0, kind: func, module: "mod", field: "func")
     OnImportFunc(import_index: 0, func_index: 0, sig_index: 0)
     OnImport(index: 1, kind: func, module: "mod", field: "func2")
@@ -39,6 +40,8 @@ BeginModule(version: 1)
     OnImportFunc(import_index: 4, func_index: 4, sig_index: 0)
     OnImport(index: 5, kind: func, module: "mod", field: "func6")
     OnImportFunc(import_index: 5, func_index: 5, sig_index: 0)
+    OnImport(index: 6, kind: func, module: "mod2", field: "func7")
+    OnImportFunc(import_index: 6, func_index: 6, sig_index: 0)
   EndImportSection
 EndModule
 ;;; STDERR ;;)
@@ -50,5 +53,6 @@ EndModule
   (import "mod" "func3" (func (;2;) (type 0)))
   (import "mod" "func4" (func (;3;) (type 0)))
   (import "mod" "func5" (func (;4;) (type 0)))
-  (import "mod" "func6" (func (;5;) (type 0))))
+  (import "mod" "func6" (func (;5;) (type 0)))
+  (import "mod2" "func7" (func (;6;) (type 0))))
 ;;; STDOUT ;;)

test/dump/compact-imports.txt

@@ -8,6 +8,7 @@
   (import "mod2" "testtable" (table 0 funcref))
   (import "mod1" "testtag" (tag (param i32)))
   (import "mod1" "func" (func (param i32 i64 f32 f64)))
+  (import "mod3" "global" (global i32))
 )
 (;; STDERR ;;;
 0000000: 0061 736d                                 ; WASM_BINARY_MAGIC
@@ -39,7 +40,7 @@
 ; section "Import" (2)
 000001b: 02                                        ; section code
 000001c: 00                                        ; section size (guess)
-000001d: 06                                        ; num imports
+000001d: 04                                        ; num imports
 ; import header 0
 000001e: 04                                        ; string length
 000001f: 6d6f 6431                                mod1  ; import module name
@@ -92,7 +93,15 @@
 000006a: 6675 6e63                                func  ; import field name
 000006e: 00                                        ; import kind
 000006f: 00                                        ; import signature index
-000001c: 53                                        ; FIXUP section size
+; import header 6
+0000070: 04                                        ; string length
+0000071: 6d6f 6433                                mod3  ; import module name
+0000075: 06                                        ; string length
+0000076: 676c 6f62 616c                           global  ; import field name
+000007c: 03                                        ; import kind
+000007d: 7f                                        ; i32
+000007e: 00                                        ; global mutability
+000001c: 62                                        ; FIXUP section size
 ;;; STDERR ;;)
 (;; STDOUT ;;;
 
@@ -104,13 +113,14 @@ Type[3]:
  - type[0] (i32, i64, f32, f64) -> nil
  - type[1] (i32) -> i32
  - type[2] (i32) -> nil
-Import[6]:
+Import[7]:
  - func[0] sig=0 <mod1.test> <- mod1.test
  - func[1] sig=1 <mod1.test2> <- mod1.test2
  - memory[0] pages: initial=0 <- mod2.testmem
  - table[0] type=funcref initial=0 <- mod2.testtable
  - tag[0] sig=2 <mod1.testtag> <- mod1.testtag
  - func[2] sig=0 <mod1.func> <- mod1.func
+ - global[0] i32 mutable=0 <- mod3.global
 
 Code Disassembly: