PCC: add a notion of memory types.

Co-authored-by: Nick Fitzgerald <fitzgen@gmail.com>

cranelift/codegen/src/ir/entities.rs

@@ -435,6 +435,8 @@ pub enum AnyEntity {
     DynamicType(DynamicType),
     /// A Global value.
     GlobalValue(GlobalValue),
+    /// A memory type.
+    MemoryType(MemoryType),
     /// A jump table.
     JumpTable(JumpTable),
     /// A constant.
@@ -460,6 +462,7 @@ impl fmt::Display for AnyEntity {
             Self::DynamicStackSlot(r) => r.fmt(f),
             Self::DynamicType(r) => r.fmt(f),
             Self::GlobalValue(r) => r.fmt(f),
+            Self::MemoryType(r) => r.fmt(f),
             Self::JumpTable(r) => r.fmt(f),
             Self::Constant(r) => r.fmt(f),
             Self::FuncRef(r) => r.fmt(f),
@@ -518,6 +521,12 @@ impl From<GlobalValue> for AnyEntity {
     }
 }
 
+impl From<MemoryType> for AnyEntity {
+    fn from(r: MemoryType) -> Self {
+        Self::MemoryType(r)
+    }
+}
+
 impl From<JumpTable> for AnyEntity {
     fn from(r: JumpTable) -> Self {
         Self::JumpTable(r)

cranelift/codegen/src/ir/function.rs

@@ -7,8 +7,8 @@ use crate::entity::{PrimaryMap, SecondaryMap};
 use crate::ir::{
     self, Block, DataFlowGraph, DynamicStackSlot, DynamicStackSlotData, DynamicStackSlots,
     DynamicType, ExtFuncData, FuncRef, GlobalValue, GlobalValueData, Inst, JumpTable,
-    JumpTableData, Layout, Opcode, SigRef, Signature, SourceLocs, StackSlot, StackSlotData,
-    StackSlots, Table, TableData, Type,
+    JumpTableData, Layout, MemoryType, MemoryTypeData, Opcode, SigRef, Signature, SourceLocs,
+    StackSlot, StackSlotData, StackSlots, Table, TableData, Type,
 };
 use crate::isa::CallConv;
 use crate::write::write_function;
@@ -173,7 +173,7 @@ pub struct FunctionStencil {
     pub global_values: PrimaryMap<ir::GlobalValue, ir::GlobalValueData>,
 
     /// Memory types for proof-carrying code.
-    pub mem_types: PrimaryMap<ir::MemoryType, ir::MemoryTypeData>,
+    pub memory_types: PrimaryMap<ir::MemoryType, ir::MemoryTypeData>,
 
     /// Tables referenced.
     pub tables: PrimaryMap<ir::Table, ir::TableData>,
@@ -204,7 +204,7 @@ impl FunctionStencil {
         self.sized_stack_slots.clear();
         self.dynamic_stack_slots.clear();
         self.global_values.clear();
-        self.mem_types.clear();
+        self.memory_types.clear();
         self.tables.clear();
         self.dfg.clear();
         self.layout.clear();
@@ -239,6 +239,11 @@ impl FunctionStencil {
         self.global_values.push(data)
     }
 
+    /// Declares a memory type for use by the function.
+    pub fn create_memory_type(&mut self, data: MemoryTypeData) -> MemoryType {
+        self.memory_types.push(data)
+    }
+
     /// Find the global dyn_scale value associated with given DynamicType.
     pub fn get_dyn_scale(&self, ty: DynamicType) -> GlobalValue {
         self.dfg.dynamic_types.get(ty).unwrap().dynamic_scale
@@ -412,7 +417,7 @@ impl Function {
                 sized_stack_slots: StackSlots::new(),
                 dynamic_stack_slots: DynamicStackSlots::new(),
                 global_values: PrimaryMap::new(),
-                mem_types: PrimaryMap::new(),
+                memory_types: PrimaryMap::new(),
                 tables: PrimaryMap::new(),
                 dfg: DataFlowGraph::new(),
                 layout: Layout::new(),

cranelift/codegen/src/ir/memtype.rs

@@ -5,6 +5,7 @@
 //! can use in proof-carrying code to validate accesses to structs and
 //! propagate facts onto the loaded values as well.
 
+use crate::ir::entities::MemoryType;
 use crate::ir::pcc::Fact;
 use crate::ir::Type;
 use alloc::vec::Vec;
@@ -15,23 +16,92 @@ use serde_derive::{Deserialize, Serialize};
 /// Data defining a memory type.
 #[derive(Clone, PartialEq, Hash)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
-pub struct MemoryTypeData {
-    /// Size of this type.
-    pub size: u64,
+pub enum MemoryTypeData {
+    /// An aggregate consisting of certain fields at certain offsets.
+    Struct {
+        /// Size of this type.
+        size: u64,
 
-    /// Fields in this type. Sorted by offset.
-    pub fields: Vec<MemoryTypeField>,
+        /// Fields in this type. Sorted by offset.
+        fields: Vec<MemoryTypeField>,
+    },
+
+    /// An aggregate consisting of a single element repeated at a
+    /// certain stride, with a statically-known length (element
+    /// count). Layout is assumed to be contiguous, with a stride
+    /// equal to the element type's size (so, if the stride is greater
+    /// than the struct's packed size, be sure to include padding in
+    /// the memory-type definition).
+    StaticArray {
+        /// The element type. May be another array, a struct, etc.
+        element: MemoryType,
+
+        /// Number of elements.
+        length: u64,
+    },
+
+    /// A single Cranelift primitive of the given type stored in
+    /// memory.
+    Primitive {
+        /// The primitive type.
+        ty: Type,
+    },
+
+    /// A type with no size.
+    Empty,
+}
+
+impl std::default::Default for MemoryTypeData {
+    fn default() -> Self {
+        Self::Empty
+    }
+}
+
+impl std::fmt::Display for MemoryTypeData {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            Self::Struct { size, fields } => {
+                write!(f, "struct {size} {{")?;
+                let mut first = true;
+                for field in fields {
+                    if first {
+                        first = false;
+                    } else {
+                        write!(f, ",")?;
+                    }
+                    write!(f, " {}: {}", field.offset, field.ty)?;
+                    if field.readonly {
+                        write!(f, " readonly")?;
+                    }
+                    if let Some(fact) = &field.fact {
+                        write!(f, " ! {}", fact)?;
+                    }
+                }
+                write!(f, " }}")?;
+                Ok(())
+            }
+            Self::StaticArray { element, length } => {
+                write!(f, "static_array {element} * {length:#x}")
+            }
+            Self::Primitive { ty } => {
+                write!(f, "primitive {ty}")
+            }
+            Self::Empty => {
+                write!(f, "empty")
+            }
+        }
+    }
 }
 
 /// One field in a memory type.
 #[derive(Clone, PartialEq, Hash)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
 pub struct MemoryTypeField {
     /// The offset of this field in the memory type.
-    pub offset: usize,
+    pub offset: u64,
     /// The type of the value in this field. Accesses to the field
     /// must use this type (i.e., cannot bitcast/type-pun in memory).
-    pub ty: Type,
+    pub ty: MemoryType,
     /// A proof-carrying-code fact about this value, if any.
     pub fact: Option<Fact>,
     /// Whether this field is read-only, i.e., stores should be

cranelift/codegen/src/ir/pcc.rs

@@ -46,6 +46,7 @@
 //! `FactContext::add()` and friends to forward-propagate facts.
 //!
 //! TODO:
+//! - Check blockparams' preds against blockparams' facts.
 //! - Propagate facts through optimization (egraph layer).
 //! - Generate facts in cranelift-wasm frontend when lowering memory ops.
 //! - Implement richer "points-to" facts that describe the pointed-to
@@ -116,6 +117,14 @@ pub enum Fact {
         /// to access (size, etc).
         region: MemoryRegion,
     },
+
+    /// A pointer to a memory type.
+    Mem {
+        /// The memory type.
+        ty: ir::MemoryType,
+        /// The offset into the memory type.
+        offset: i64,
+    },
 }
 
 /// A memory region that can be accessed. This description is attached
@@ -132,10 +141,11 @@ pub struct MemoryRegion {
 impl fmt::Display for Fact {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         match self {
-            Fact::ValueMax { bit_width, max } => write!(f, "max({}, 0x{:x})", bit_width, max),
+            Fact::ValueMax { bit_width, max } => write!(f, "max({}, {:#x})", bit_width, max),
             Fact::PointsTo {
                 region: MemoryRegion { max },
             } => write!(f, "points_to(0x{:x})", max),
+            Fact::Mem { ty, offset } => write!(f, "mem({}, {:#x})", ty, offset),
         }
     }
 }

cranelift/codegen/src/isa/aarch64/pcc.rs

@@ -28,6 +28,12 @@ fn check_subsumes(ctx: &FactContext, subsumer: &Fact, subsumee: &Fact) -> PccRes
         subsumer,
         subsumee
     );
+
+    // For now, allow all `mem` facts to validate.
+    if matches!(subsumee, Fact::Mem { .. }) {
+        return Ok(());
+    }
+
     if ctx.subsumes(subsumer, subsumee) {
         Ok(())
     } else {

cranelift/codegen/src/write.rs

@@ -56,6 +56,11 @@ pub trait FuncWriter {
             self.write_entity_definition(w, func, gv.into(), gv_data)?;
         }
 
+        for (mt, mt_data) in &func.memory_types {
+            any = true;
+            self.write_entity_definition(w, func, mt.into(), mt_data)?;
+        }
+
         for (table, table_data) in &func.tables {
             if !table_data.index_type.is_invalid() {
                 any = true;

cranelift/filetests/filetests/pcc/succeed/memtypes.clif

@@ -3,20 +3,28 @@ set enable_pcc=true
 target aarch64
 
 function %f0(i64) -> i32 {
-    mt0 = memtype 8 { 0: i32, 4: i32 readonly }
+    mt0 = struct 8 { 0: mt1, 4: mt1 readonly }
+    mt1 = primitive i32
 
-block0(v0 ! memtype(mt0, 0): i64):  ;; v0 points to an instance of mt0, at offset 0
-    v1 = load.i32 checked v0+0
-    v2 = load.i32 checked v0+4
+block0(v0 ! mem(mt0, 0): i64):  ;; v0 points to an instance of mt0, at offset 0
+    v1 = load.i32 v0+0
+    v2 = load.i32 v0+4
     v3 = iadd.i32 v1, v2
     return v3
 }
 
 function %f1(i64) -> i32 {
-    mt0 = memtype 8 { 0: i64 readonly ! points_to(0x1_0000_0000 }
+    ;; TODO: validate that fields don't overlap, and are within the
+    ;; overall size. Make note of future thoughts for enums, tag
+    ;; discriminants, etc.
+    mt0 = struct 8 { 0: mt3 readonly ! mem(mt1, 0) }
+    ;; `u8` here could be a CLIF type or a memtype
+    mt1 = static_array mt2 * 0x1_0000_0000
+    mt2 = primitive i8
+    mt3 = primitive i64
 
-block0(v0 ! memtype(mt0, 0): i64):
-    v1 ! points_to(0x1_0000_0000) = load.i64 checked v0
-    v2 = load.i32 checked v1+0x1000
+block0(v0 ! mem(mt0, 0): i64):
+    v1 ! mem(mt1, 0) = load.i64 v0
+    v2 = load.i32 v1+0x1000
     return v2
 }

cranelift/reader/src/lexer.rs

@@ -40,6 +40,7 @@ pub enum Token<'a> {
     StackSlot(u32),        // ss3
     DynamicStackSlot(u32), // dss4
     GlobalValue(u32),      // gv3
+    MemoryType(u32),       // mt0
     Table(u32),            // table2
     Constant(u32),         // const2
     FuncRef(u32),          // fn2
@@ -344,6 +345,7 @@ impl<'a> Lexer<'a> {
             "dss" => Some(Token::DynamicStackSlot(number)),
             "dt" => Some(Token::DynamicType(number)),
             "gv" => Some(Token::GlobalValue(number)),
+            "mt" => Some(Token::MemoryType(number)),
             "table" => Some(Token::Table(number)),
             "const" => Some(Token::Constant(number)),
             "fn" => Some(Token::FuncRef(number)),