Use precise alias regions for statically-known entity imports

crates/cranelift/src/alias_region.rs

@@ -75,8 +75,8 @@ enum AliasRegionKey {
         offset: u32,
     },
 
-    /// An imported or exported memory access (shared across all
-    /// imported/exported memories).
+    /// An access of a memory that crosses a module boundary and whose
+    /// definition we do not statically know (shared across all such memories).
     PublicMemory,
 
     /// A defined memory access.
@@ -87,8 +87,8 @@ enum AliasRegionKey {
         index: DefinedMemoryIndex,
     },
 
-    /// An imported or exported table access (shared across all
-    /// imported/exported tables).
+    /// An access of a table that crosses a module boundary and whose definition
+    /// we do not statically know (shared across all such tables).
     PublicTable,
 
     /// A defined table access.
@@ -99,8 +99,8 @@ enum AliasRegionKey {
         index: DefinedTableIndex,
     },
 
-    /// An imported or exported global access (shared across all
-    /// imported/exported globals).
+    /// An access of a global that crosses a module boundary and whose definition
+    /// we do not statically know (shared across all such globals).
     PublicGlobal,
 
     /// A defined global access.
@@ -937,14 +937,13 @@ where
         self.region(func, AliasRegionKey::GcHeap)
     }
 
-    /// Get the alias region for an imported or exported memory access (shared
-    /// across all imported/exported memories).
+    /// Get the alias region shared by all memories that cross a module boundary
+    /// and whose definition we do not statically know.
     pub fn public_memory_region(&mut self, func: &mut ir::Function) -> ir::AliasRegion {
         self.region(func, AliasRegionKey::PublicMemory)
     }
 
-    /// Get the alias region for accessing a defined memory that is not
-    /// exported.
+    /// Get the alias region for accessing a particular defined memory.
     pub fn defined_memory_region(
         &mut self,
         func: &mut ir::Function,
@@ -954,14 +953,13 @@ where
         self.region(func, AliasRegionKey::DefinedMemory { module, index })
     }
 
-    /// Get the alias region for an imported or exported table access (shared
-    /// across all imported/exported memories).
+    /// Get the alias region shared by all tables that cross a module boundary
+    /// and whose definition we do not statically know.
     pub fn public_table_region(&mut self, func: &mut ir::Function) -> ir::AliasRegion {
         self.region(func, AliasRegionKey::PublicTable)
     }
 
-    /// Get the alias region for accessing a defined table that is not
-    /// exported.
+    /// Get the alias region for accessing a particular defined table.
     pub fn defined_table_region(
         &mut self,
         func: &mut ir::Function,
@@ -971,14 +969,13 @@ where
         self.region(func, AliasRegionKey::DefinedTable { module, index })
     }
 
-    /// Get the alias region for an imported or exported global access (shared
-    /// across all imported/exported memories).
+    /// Get the alias region shared by all globals that cross a module boundary
+    /// and whose definition we do not statically know.
     pub fn public_global_region(&mut self, func: &mut ir::Function) -> ir::AliasRegion {
         self.region(func, AliasRegionKey::PublicGlobal)
     }
 
-    /// Get the alias region for accessing a defined global that is not
-    /// exported.
+    /// Get the alias region for accessing a particular defined global.
     pub fn defined_global_region(
         &mut self,
         func: &mut ir::Function,

crates/cranelift/src/func_environ.rs

@@ -342,70 +342,116 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
         self.isa.pointer_type()
     }
 
+    /// Get the alias region to use for accesses of the given memory.
+    ///
+    /// XXX: Keep the `{memory,global,table}_alias_region` methods in sync with
+    /// each other.
     pub(crate) fn memory_alias_region(
         &mut self,
         func: &mut Function,
         memory: MemoryIndex,
     ) -> ir::AliasRegion {
-        if self.module.is_exported_memory(memory) {
-            // A function that operates on an exported defined memory can be
-            // inlined into a different module caller, where that that caller's
-            // module also imports that exported memory. That caller will access
-            // the memory with `AliasRegionKey::PublicMemory`, so we must also
-            // conservatively do the same here, even though we potentially know
-            // the precise static module index and defined memory index, because
-            // memory accessed with two different alias regions must not
-            // actually alias, or else we will get miscompiles.
-            self.alias_regions.public_memory_region(func)
-        } else {
-            match self.module.defined_memory_index(memory) {
-                Some(def) => self.alias_regions.defined_memory_region(
-                    func,
-                    self.translation.module_index(),
-                    def,
-                ),
-                None => self.alias_regions.public_memory_region(func),
+        match self.module.defined_memory_index(memory) {
+            // A memory defined by this module. When it is exported, a function
+            // that operates on it can be inlined into a caller in a different
+            // module that imports that memory, and vice versa. That other module
+            // accesses the memory with `AliasRegionKey::PublicMemory` unless it
+            // statically knows that its import is always this memory, so we can
+            // only use this memory's precise region when every module that may
+            // import it does know that. Memory accessed with two different alias
+            // regions must not actually alias, or else we will get miscompiles.
+            Some(def) => {
+                if self.module.is_exported_memory(memory)
+                    && !self.translation.memories_known_to_importers.contains(def)
+                {
+                    self.alias_regions.public_memory_region(func)
+                } else {
+                    self.alias_regions.defined_memory_region(
+                        func,
+                        self.translation.module_index(),
+                        def,
+                    )
+                }
             }
+
+            // A memory imported by this module: use the precise region when we
+            // statically know which defined memory always satisfies the import
+            // and everything else that imports it knows the same.
+            None => match self.translation.known_imported_memories[memory] {
+                Some(known) => {
+                    self.alias_regions
+                        .defined_memory_region(func, known.module, known.index)
+                }
+                None => self.alias_regions.public_memory_region(func),
+            },
         }
     }
 
+    /// Get the alias region to use for accesses of the given table.
+    ///
+    /// XXX: Keep the `{memory,global,table}_alias_region` methods in sync with
+    /// each other.
     pub(crate) fn table_alias_region(
         &mut self,
         func: &mut Function,
         table: TableIndex,
     ) -> ir::AliasRegion {
-        if self.module.is_exported_table(table) {
-            // See the comment in `memory_alias_region` for details.
-            self.alias_regions.public_table_region(func)
-        } else {
-            match self.module.defined_table_index(table) {
-                Some(def) => self.alias_regions.defined_table_region(
-                    func,
-                    self.translation.module_index(),
-                    def,
-                ),
-                None => self.alias_regions.public_table_region(func),
+        // See the comments in `memory_alias_region` for details.
+        match self.module.defined_table_index(table) {
+            Some(def) => {
+                if self.module.is_exported_table(table)
+                    && !self.translation.tables_known_to_importers.contains(def)
+                {
+                    self.alias_regions.public_table_region(func)
+                } else {
+                    self.alias_regions.defined_table_region(
+                        func,
+                        self.translation.module_index(),
+                        def,
+                    )
+                }
             }
+            None => match self.translation.known_imported_tables[table] {
+                Some(known) => {
+                    self.alias_regions
+                        .defined_table_region(func, known.module, known.index)
+                }
+                None => self.alias_regions.public_table_region(func),
+            },
         }
     }
 
+    /// Get the alias region to use for accesses of the given global.
+    ///
+    /// XXX: Keep the `{memory,global,table}_alias_region` methods in sync with
+    /// each other.
     pub(crate) fn global_alias_region(
         &mut self,
         func: &mut Function,
         global: GlobalIndex,
     ) -> ir::AliasRegion {
-        if self.module.is_exported_global(global) {
-            // See the comment in `memory_alias_region` for details.
-            self.alias_regions.public_global_region(func)
-        } else {
-            match self.module.defined_global_index(global) {
-                Some(def) => self.alias_regions.defined_global_region(
-                    func,
-                    self.translation.module_index(),
-                    def,
-                ),
-                None => self.alias_regions.public_global_region(func),
+        // See the comments in `memory_alias_region` for details.
+        match self.module.defined_global_index(global) {
+            Some(def) => {
+                if self.module.is_exported_global(global)
+                    && !self.translation.globals_known_to_importers.contains(def)
+                {
+                    self.alias_regions.public_global_region(func)
+                } else {
+                    self.alias_regions.defined_global_region(
+                        func,
+                        self.translation.module_index(),
+                        def,
+                    )
+                }
             }
+            None => match self.translation.known_imported_globals[global] {
+                Some(known) => {
+                    self.alias_regions
+                        .defined_global_region(func, known.module, known.index)
+                }
+                None => self.alias_regions.public_global_region(func),
+            },
         }
     }
 

crates/environ/src/collections/entity_set.rs

@@ -3,14 +3,25 @@ use wasmtime_core::error::OutOfMemory;
 
 /// Like `cranelift_entity::EntitySet` but enforces fallible allocation for all
 /// methods that allocate.
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct TryEntitySet<K>
 where
     K: EntityRef,
 {
     inner: cranelift_entity::EntitySet<K>,
 }
 
+impl<K> Default for TryEntitySet<K>
+where
+    K: EntityRef,
+{
+    fn default() -> Self {
+        Self {
+            inner: Default::default(),
+        }
+    }
+}
+
 impl<K> TryEntitySet<K>
 where
     K: EntityRef,

crates/environ/src/compile/module_environ.rs

@@ -4,12 +4,13 @@ use crate::module::{
 };
 use crate::prelude::*;
 use crate::{
-    ConstExpr, ConstOp, DataIndex, DefinedFuncIndex, DefinedGlobalIndex, ElemIndex,
-    EngineOrModuleTypeIndex, EntityIndex, EntityType, FuncIndex, FuncKey, GlobalIndex, IndexType,
-    MemoryIndex, MemoryInitializer, ModuleInternedTypeIndex, ModuleStartup, ModuleTypesBuilder,
-    PanicOnOom as _, PassiveElemIndex, PrimaryMap, RuntimeDataIndex, StaticModuleIndex, TableIndex,
-    TableInitialValue, TableInitialization, Tag, TagIndex, Trap, Tunables, TypeConvert, TypeIndex,
-    WasmHeapTopType, WasmHeapType, WasmResult, WasmValType, WasmparserTypeConverter,
+    ConstExpr, ConstOp, DataIndex, DefinedFuncIndex, DefinedGlobalIndex, DefinedMemoryIndex,
+    DefinedTableIndex, ElemIndex, EngineOrModuleTypeIndex, EntityIndex, EntityType, FuncIndex,
+    FuncKey, GlobalIndex, IndexType, MemoryIndex, MemoryInitializer, ModuleInternedTypeIndex,
+    ModuleStartup, ModuleTypesBuilder, PanicOnOom as _, PassiveElemIndex, PrimaryMap,
+    RuntimeDataIndex, StaticModuleIndex, TableIndex, TableInitialValue, TableInitialization, Tag,
+    TagIndex, Trap, Tunables, TypeConvert, TypeIndex, WasmHeapTopType, WasmHeapType, WasmResult,
+    WasmValType, WasmparserTypeConverter,
 };
 use alloc::borrow::Cow;
 use cranelift_entity::SecondaryMap;
@@ -72,6 +73,15 @@ impl From<FactInlineIntrinsic> for KnownFunc {
     }
 }
 
+/// A statically-known import of a core Wasm global, memory, or table.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub struct KnownEntity<T> {
+    /// The module that defines this entity.
+    pub module: StaticModuleIndex,
+    /// The entity's index in the defining module's defined-entity index space.
+    pub index: T,
+}
+
 /// The result of translating via `ModuleEnvironment`.
 ///
 /// Function bodies are not yet translated, and data initializers have not yet
@@ -107,6 +117,47 @@ pub struct ModuleTranslation<'data> {
     /// `FuncKey::FactInlineIntrinsic`s.
     pub known_imported_functions: SecondaryMap<FuncIndex, Option<KnownFunc>>,
 
+    /// For each imported global, memory, or table, the single statically-known
+    /// defined entity that always satisfies that import, if any.
+    ///
+    /// This is used to access the entity via the defining module's precise
+    /// `AliasRegionKey::Defined{Global,Memory,Table}` region instead of the
+    /// conservative `AliasRegionKey::Public{Global,Memory,Table}` region that is
+    /// shared by every entity of that kind which crosses a module boundary.
+    ///
+    /// XXX: Being known requires more here than it does for functions: it is
+    /// not enough that *this* module's import is always the same entity,
+    /// *every* module that may import that entity must also always import that
+    /// same entity. Otherwise a function from one of those other modules, which
+    /// accesses the entity via the conservative region, could be inlined next
+    /// to one of our accesses via the precise region, and accessing the same
+    /// memory through two different alias regions is invalid.
+    pub known_imported_globals: SecondaryMap<GlobalIndex, Option<KnownEntity<DefinedGlobalIndex>>>,
+
+    /// Same as `known_imported_globals`, but for memories.
+    pub known_imported_memories: SecondaryMap<MemoryIndex, Option<KnownEntity<DefinedMemoryIndex>>>,
+
+    /// Same as `known_imported_globals`, but for tables.
+    pub known_imported_tables: SecondaryMap<TableIndex, Option<KnownEntity<DefinedTableIndex>>>,
+
+    /// For each global defined by this module, whether every module that may
+    /// import this global always imports exactly this global.
+    ///
+    /// When this holds, accesses of the global may use its precise
+    /// `AliasRegionKey::DefinedGlobal` region even when the global is exported,
+    /// because every module that can reach it agrees on that same region. This
+    /// is vacuously true of globals that nothing in the component imports.
+    ///
+    /// This can only be determined by looking at the whole component, so it is
+    /// always `false` for standalone modules.
+    pub globals_known_to_importers: TryEntitySet<DefinedGlobalIndex>,
+
+    /// Same as [`Self::globals_known_to_importers`], but for memories.
+    pub memories_known_to_importers: TryEntitySet<DefinedMemoryIndex>,
+
+    /// Same as [`Self::globals_known_to_importers`], but for tables.
+    pub tables_known_to_importers: TryEntitySet<DefinedTableIndex>,
+
     /// A list of type signatures which are considered exported from this
     /// module, or those that can possibly be called. This list is sorted, and
     /// trampolines for each of these signatures are required.
@@ -228,6 +279,12 @@ impl<'data> ModuleTranslation<'data> {
             wasm_module_offset: 0,
             function_body_inputs: PrimaryMap::default(),
             known_imported_functions: SecondaryMap::default(),
+            known_imported_globals: SecondaryMap::default(),
+            known_imported_memories: SecondaryMap::default(),
+            known_imported_tables: SecondaryMap::default(),
+            globals_known_to_importers: TryEntitySet::default(),
+            memories_known_to_importers: TryEntitySet::default(),
+            tables_known_to_importers: TryEntitySet::default(),
             exported_signatures: Vec::default(),
             debuginfo: DebugInfoData::default(),
             has_unparsed_debuginfo: false,