Add ConstantPool

cranelift/codegen/src/ir/constant.rs

@@ -0,0 +1,218 @@
+//! Constants
+//!
+//! The constant pool defined here allows cranelift to avoid emitting the same constant multiple
+//! times. As constants are inserted in the pool, a handle is returned; the handle is a cranelift
+//! Entity. Inserting the same data multiple times will always return the same handle. Future work
+//! could include: ensuring alignment of constants within the pool, bucketing constants by size.
+
+use crate::ir::Constant;
+use cranelift_entity::EntityRef;
+use std::collections::{BTreeMap, HashMap};
+use std::vec::Vec;
+
+/// This type describes the actual constant data.
+pub type ConstantData = Vec<u8>;
+
+/// This type describes an offset in bytes within a constant pool.
+pub type ConstantOffset = u32;
+
+/// Inner type for storing data and offset together in the constant pool. The offset is optional
+/// because it must be set relative to the function code size (i.e. constants are emitted after the
+/// function body); because the function is not yet compiled when constants are inserted,
+/// [set_offset](ir::ConstantPool::set_offset) must be called once a constant's offset from the
+/// beginning of the function is known (see [relaxation.rs](binemit::relaxation)).
+#[derive(Clone)]
+pub struct ConstantPoolEntry {
+    data: ConstantData,
+    offset: Option<ConstantOffset>,
+}
+
+impl ConstantPoolEntry {
+    fn new(data: ConstantData) -> Self {
+        ConstantPoolEntry { data, offset: None }
+    }
+
+    /// Return the size of the constant at this entry.
+    pub fn len(&self) -> usize {
+        self.data.len()
+    }
+
+    /// Assign a new offset to the constant at this entry.
+    pub fn set_offset(&mut self, offset: ConstantOffset) {
+        self.offset = Some(offset)
+    }
+}
+
+/// Maintains the mapping between a constant handle (i.e. [Constant](ir::entities::Constant)) and
+/// its constant data (i.e. [ConstantData](ir::constant::ConstantData)).
+#[derive(Clone)]
+pub struct ConstantPool {
+    /// This mapping maintains the insertion order as long as Constants are created with sequentially increasing integers.
+    handles_to_values: BTreeMap<Constant, ConstantPoolEntry>,
+    /// This mapping is unordered (no need for lexicographic ordering) but allows us to map constant data back to handles.
+    values_to_handles: HashMap<ConstantData, Constant>,
+}
+
+impl ConstantPool {
+    /// Create a new constant pool instance.
+    pub fn new() -> Self {
+        ConstantPool {
+            handles_to_values: BTreeMap::new(),
+            values_to_handles: HashMap::new(),
+        }
+    }
+
+    /// Empty the constant pool of all data.
+    pub fn clear(&mut self) {
+        self.handles_to_values.clear();
+        self.values_to_handles.clear();
+    }
+
+    /// Insert constant data into the pool, returning a handle for later referencing; when constant
+    /// data is inserted that is a duplicate of previous constant data, the existing handle will be
+    /// returned.
+    pub fn insert(&mut self, constant_value: ConstantData) -> Constant {
+        if self.values_to_handles.contains_key(&constant_value) {
+            self.values_to_handles
+                .get(&constant_value)
+                .expect("A constant handle must have a corresponding constant value; this is an implementation error in ConstantPool")
+                .clone()
+        } else {
+            let constant_handle = Constant::new(self.len());
+            self.values_to_handles
+                .insert(constant_value.clone(), constant_handle.clone());
+            self.handles_to_values.insert(
+                constant_handle.clone(),
+                ConstantPoolEntry::new(constant_value),
+            );
+            constant_handle
+        }
+    }
+
+    /// Retrieve the constant data given a handle.
+    pub fn get(&self, constant_handle: Constant) -> &ConstantData {
+        assert!(self.handles_to_values.contains_key(&constant_handle));
+        &self.handles_to_values
+            .get(&constant_handle)
+            .expect("A constant handle must have a corresponding constant value; was a constant handle created outside of the pool?")
+            .data
+    }
+
+    /// Assign an offset to a given constant, where the offset is the number of bytes from the
+    /// beginning of the function to the beginning of the constant data inside the pool.
+    pub fn set_offset(&mut self, constant_handle: Constant, constant_offset: ConstantOffset) {
+        assert!(self.handles_to_values.contains_key(&constant_handle), "A constant handle must have already been inserted into the pool; perhaps a constant pool was created outside of the pool?");
+        self.handles_to_values
+            .entry(constant_handle)
+            .and_modify(|e| e.offset = Some(constant_offset));
+    }
+
+    /// Retrieve the offset of a given constant, where the offset is the number of bytes from the
+    /// beginning of the function to the beginning of the constant data inside the pool.
+    pub fn get_offset(&self, constant_handle: Constant) -> ConstantOffset {
+        self.handles_to_values.get(&constant_handle)
+            .expect("A constant handle must have a corresponding constant value; was a constant handle created outside of the pool?")
+            .offset
+            .expect("A constant offset has not yet been set; verify that `set_offset` has been called before this point")
+    }
+
+    /// Iterate over the constants in insertion order.
+    pub fn iter(&self) -> impl Iterator<Item = (&Constant, &ConstantData)> {
+        self.handles_to_values.iter().map(|(h, e)| (h, &e.data))
+    }
+
+    /// Iterate over mutable entries in the constant pool in insertion order.
+    pub fn entries_mut(&mut self) -> impl Iterator<Item = &mut ConstantPoolEntry> {
+        self.handles_to_values.values_mut()
+    }
+
+    /// Return the number of constants in the pool.
+    pub fn len(&self) -> usize {
+        self.handles_to_values.len()
+    }
+
+    /// Return the combined size of all of the constant values in the pool.
+    pub fn byte_size(&self) -> usize {
+        self.values_to_handles.keys().map(|c| c.len()).sum()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn empty() {
+        let sut = ConstantPool::new();
+        assert_eq!(sut.len(), 0);
+    }
+
+    #[test]
+    fn insert() {
+        let mut sut = ConstantPool::new();
+        sut.insert(vec![1, 2, 3]);
+        sut.insert(vec![4, 5, 6]);
+        assert_eq!(sut.len(), 2);
+    }
+
+    #[test]
+    fn insert_duplicate() {
+        let mut sut = ConstantPool::new();
+        let a = sut.insert(vec![1, 2, 3]);
+        sut.insert(vec![4, 5, 6]);
+        let b = sut.insert(vec![1, 2, 3]);
+        assert_eq!(a, b);
+    }
+
+    #[test]
+    fn clear() {
+        let mut sut = ConstantPool::new();
+        sut.insert(vec![1, 2, 3]);
+        assert_eq!(sut.len(), 1);
+
+        sut.clear();
+        assert_eq!(sut.len(), 0);
+    }
+
+    #[test]
+    fn iteration_order() {
+        let mut sut = ConstantPool::new();
+        sut.insert(vec![1, 2, 3]);
+        sut.insert(vec![4, 5, 6]);
+        sut.insert(vec![1, 2, 3]);
+        let data = sut.iter().map(|(_, v)| v).collect::<Vec<&ConstantData>>();
+        assert_eq!(data, vec![&vec![1, 2, 3], &vec![4, 5, 6]]);
+    }
+
+    #[test]
+    fn get() {
+        let mut sut = ConstantPool::new();
+        let data = vec![1, 2, 3];
+        let handle = sut.insert(data.clone());
+        assert_eq!(sut.get(handle), &data);
+    }
+
+    #[test]
+    #[should_panic]
+    fn get_nonexistent_constant() {
+        let sut = ConstantPool::new();
+        let a = Constant::with_number(42).unwrap();
+        sut.get(a); // panics, only use constants returned by ConstantPool
+    }
+
+    #[test]
+    fn get_offset() {
+        let mut sut = ConstantPool::new();
+        let a = sut.insert(vec![1]);
+        sut.set_offset(a, 42);
+        assert_eq!(sut.get_offset(a), 42)
+    }
+
+    #[test]
+    #[should_panic]
+    fn get_nonexistent_offset() {
+        let mut sut = ConstantPool::new();
+        let a = sut.insert(vec![1]);
+        sut.get_offset(a); // panics, set_offset should have been called
+    }
+}

cranelift/codegen/src/ir/dfg.rs

@@ -5,7 +5,7 @@ use crate::ir;
 use crate::ir::builder::ReplaceBuilder;
 use crate::ir::extfunc::ExtFuncData;
 use crate::ir::instructions::{BranchInfo, CallInfo, InstructionData};
-use crate::ir::types;
+use crate::ir::{types, ConstantPool};
 use crate::ir::{
     Ebb, FuncRef, Inst, SigRef, Signature, Type, Value, ValueLabelAssignments, ValueList,
     ValueListPool,
@@ -67,6 +67,9 @@ pub struct DataFlowGraph {
 
     /// Saves Value labels.
     pub values_labels: Option<HashMap<Value, ValueLabelAssignments>>,
+
+    /// Constants used within the function
+    pub constants: ConstantPool,
 }
 
 impl DataFlowGraph {
@@ -81,6 +84,7 @@ impl DataFlowGraph {
             signatures: PrimaryMap::new(),
             ext_funcs: PrimaryMap::new(),
             values_labels: None,
+            constants: ConstantPool::new(),
         }
     }
 
@@ -94,6 +98,7 @@ impl DataFlowGraph {
         self.signatures.clear();
         self.ext_funcs.clear();
         self.values_labels = None;
+        self.constants.clear()
     }
 
     /// Get the total number of instructions created in this function, whether they are currently

cranelift/codegen/src/ir/entities.rs

@@ -104,6 +104,24 @@ impl GlobalValue {
     }
 }
 
+/// An opaque reference to a constant
+#[derive(Copy, Clone, PartialEq, Eq, Hash, Ord, PartialOrd)]
+pub struct Constant(u32);
+entity_impl!(Constant, "const");
+
+impl Constant {
+    /// Create a const reference from its number.
+    ///
+    /// This method is for use by the parser.
+    pub fn with_number(n: u32) -> Option<Self> {
+        if n < u32::MAX {
+            Some(Constant(n))
+        } else {
+            None
+        }
+    }
+}
+
 /// An opaque reference to a jump table.
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 #[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
@@ -195,7 +213,7 @@ impl Table {
     }
 }
 
-/// A reference to any of the entities defined in this module.
+/// A reference to any of the entities defined in this module that can appear in CLIF IR.
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub enum AnyEntity {
     /// The whole function.
@@ -331,4 +349,10 @@ mod tests {
             mem::size_of::<PackedOption<Value>>()
         );
     }
+
+    #[test]
+    fn constant_with_number() {
+        assert_eq!(Constant::with_number(0).unwrap().to_string(), "const0");
+        assert_eq!(Constant::with_number(1).unwrap().to_string(), "const1");
+    }
 }

cranelift/codegen/src/ir/mod.rs

@@ -2,6 +2,7 @@
 
 mod builder;
 pub mod condcodes;
+pub mod constant;
 pub mod dfg;
 pub mod entities;
 mod extfunc;
@@ -27,9 +28,10 @@ mod valueloc;
 use serde::{Deserialize, Serialize};
 
 pub use crate::ir::builder::{InsertBuilder, InstBuilder, InstBuilderBase, InstInserterBase};
+pub use crate::ir::constant::{ConstantData, ConstantOffset, ConstantPool};
 pub use crate::ir::dfg::{DataFlowGraph, ValueDef};
 pub use crate::ir::entities::{
-    Ebb, FuncRef, GlobalValue, Heap, Inst, JumpTable, SigRef, StackSlot, Table, Value,
+    Constant, Ebb, FuncRef, GlobalValue, Heap, Inst, JumpTable, SigRef, StackSlot, Table, Value,
 };
 pub use crate::ir::extfunc::{
     AbiParam, ArgumentExtension, ArgumentPurpose, ExtFuncData, Signature,