Allow creating constants in legalization AST

This adds a `DummyConstant` structure that is converted to something like `let const0 = pos.func.dfg.constants.insert(...)` in `gen_legalizer.rs`. This allows us to create constants during legalization with something like `let ones = constant(vec![0xff; 16])` and then use `ones` within a `def!` block, e.g.: `def!(a = vconst(ones))`. One unfortunate side-effect of this change is that, because the names of the constants in `ConstPool` are dynamic, the `VarPool` and `SymbolTable` structures that previously operated on `&'static str` types now must operate on `String` types; however, since this is a change to the meta code-generation, it should result in no runtime performance impact.

cranelift/codegen/meta/src/cdsl/ast.rs

@@ -54,7 +54,7 @@ impl Def {
         let results = self
             .defined_vars
             .iter()
-            .map(|&x| var_pool.get(x).name)
+            .map(|&x| var_pool.get(x).name.as_str())
             .collect::<Vec<_>>();
 
         let results = if results.len() == 1 {
@@ -238,7 +238,7 @@ pub(crate) enum PatternPosition {
 ///
 /// Temporary values are defined only in the destination pattern.
 pub(crate) struct Var {
-    pub name: &'static str,
+    pub name: String,
 
     /// The `Def` defining this variable in a source pattern.
     pub src_def: Option<DefIndex>,
@@ -254,7 +254,7 @@ pub(crate) struct Var {
 }
 
 impl Var {
-    fn new(name: &'static str) -> Self {
+    fn new(name: String) -> Self {
         Self {
             name,
             src_def: None,
@@ -346,7 +346,7 @@ impl Var {
     }
 
     pub fn to_rust_code(&self) -> String {
-        self.name.into()
+        self.name.clone()
     }
     fn rust_type(&self) -> String {
         self.type_var.as_ref().unwrap().to_rust_code()
@@ -384,8 +384,51 @@ impl VarPool {
     pub fn get_mut(&mut self, index: VarIndex) -> &mut Var {
         self.pool.get_mut(index).unwrap()
     }
-    pub fn create(&mut self, name: &'static str) -> VarIndex {
+    pub fn create(&mut self, name: impl Into<String>) -> VarIndex {
-        self.pool.push(Var::new(name))
+        self.pool.push(Var::new(name.into()))
+    }
+}
+
+/// Contains constants created in the AST that must be inserted into the true [ConstantPool]
+/// (cranelift_codegen::ir::ConstantPool) when the legalizer code is generated. The constant data
+/// is named in the order it is inserted; inserting data using [insert]
+/// (cranelift_codegen_meta::cdsl::ast::insert) will avoid duplicates.
+pub(crate) struct ConstPool {
+    pool: Vec<Vec<u8>>,
+}
+
+impl ConstPool {
+    /// Create an empty constant pool.
+    pub fn new() -> Self {
+        Self { pool: vec![] }
+    }
+
+    /// Create a name for a constant from its position in the pool.
+    fn create_name(position: usize) -> String {
+        format!("const{}", position)
+    }
+
+    /// Insert constant data into the pool, returning the name of the variable used to reference it.
+    /// This method will search for data that matches the new data and return the existing constant
+    /// name to avoid duplicates.
+    pub fn insert(&mut self, data: Vec<u8>) -> String {
+        let possible_position = self.pool.iter().position(|d| d == &data);
+        let position = if let Some(found_position) = possible_position {
+            found_position
+        } else {
+            let new_position = self.pool.len();
+            self.pool.push(data);
+            new_position
+        };
+        ConstPool::create_name(position)
+    }
+
+    /// Iterate over the name/value pairs in the pool.
+    pub fn iter(&self) -> impl Iterator<Item = (String, &Vec<u8>)> {
+        self.pool
+            .iter()
+            .enumerate()
+            .map(|(i, v)| (ConstPool::create_name(i), v))
     }
 }
 
@@ -567,13 +610,14 @@ impl Apply {
 pub(crate) enum DummyExpr {
     Var(DummyVar),
     Literal(Literal),
+    Constant(DummyConstant),
     Apply(InstSpec, Vec<DummyExpr>),
     Block(DummyVar),
 }
 
 #[derive(Clone)]
 pub(crate) struct DummyVar {
-    pub name: &'static str,
+    pub name: String,
 }
 
 impl Into<DummyExpr> for DummyVar {
@@ -587,8 +631,23 @@ impl Into<DummyExpr> for Literal {
     }
 }
 
-pub(crate) fn var(name: &'static str) -> DummyVar {
+#[derive(Clone)]
-    DummyVar { name }
+pub(crate) struct DummyConstant(pub(crate) Vec<u8>);
+
+pub(crate) fn constant(data: Vec<u8>) -> DummyConstant {
+    DummyConstant(data)
+}
+
+impl Into<DummyExpr> for DummyConstant {
+    fn into(self) -> DummyExpr {
+        DummyExpr::Constant(self)
+    }
+}
+
+pub(crate) fn var(name: &str) -> DummyVar {
+    DummyVar {
+        name: name.to_owned(),
+    }
 }
 
 pub(crate) struct DummyDef {
@@ -656,3 +715,40 @@ macro_rules! ebb {
         ExprBuilder::block($block).assign_to(Vec::new())
     };
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::cdsl::ast::ConstPool;
+
+    #[test]
+    fn const_pool_returns_var_names() {
+        let mut c = ConstPool::new();
+        assert_eq!(c.insert([0, 1, 2].to_vec()), "const0");
+        assert_eq!(c.insert([1, 2, 3].to_vec()), "const1");
+    }
+
+    #[test]
+    fn const_pool_avoids_duplicates() {
+        let data = [0, 1, 2].to_vec();
+        let mut c = ConstPool::new();
+        assert_eq!(c.pool.len(), 0);
+
+        assert_eq!(c.insert(data.clone()), "const0");
+        assert_eq!(c.pool.len(), 1);
+
+        assert_eq!(c.insert(data), "const0");
+        assert_eq!(c.pool.len(), 1);
+    }
+
+    #[test]
+    fn const_pool_iterates() {
+        let mut c = ConstPool::new();
+        c.insert([0, 1, 2].to_vec());
+        c.insert([3, 4, 5].to_vec());
+
+        let mut iter = c.iter();
+        assert_eq!(iter.next(), Some(("const0".to_owned(), &vec![0, 1, 2])));
+        assert_eq!(iter.next(), Some(("const1".to_owned(), &vec![3, 4, 5])));
+        assert_eq!(iter.next(), None);
+    }
+}

cranelift/codegen/meta/src/cdsl/xform.rs

@@ -1,6 +1,6 @@
 use crate::cdsl::ast::{
-    Apply, BlockPool, DefIndex, DefPool, DummyDef, DummyExpr, Expr, PatternPosition, VarIndex,
+    Apply, BlockPool, ConstPool, DefIndex, DefPool, DummyDef, DummyExpr, Expr, PatternPosition,
-    VarPool,
+    VarIndex, VarPool,
 };
 use crate::cdsl::instructions::Instruction;
 use crate::cdsl::type_inference::{infer_transform, TypeEnvironment};
@@ -24,16 +24,18 @@ pub(crate) struct Transform {
     pub var_pool: VarPool,
     pub def_pool: DefPool,
     pub block_pool: BlockPool,
+    pub const_pool: ConstPool,
     pub type_env: TypeEnvironment,
 }
 
-type SymbolTable = HashMap<&'static str, VarIndex>;
+type SymbolTable = HashMap<String, VarIndex>;
 
 impl Transform {
     fn new(src: DummyDef, dst: Vec<DummyDef>) -> Self {
         let mut var_pool = VarPool::new();
         let mut def_pool = DefPool::new();
         let mut block_pool = BlockPool::new();
+        let mut const_pool = ConstPool::new();
 
         let mut input_vars: Vec<VarIndex> = Vec::new();
         let mut defined_vars: Vec<VarIndex> = Vec::new();
@@ -51,6 +53,7 @@ impl Transform {
             &mut var_pool,
             &mut def_pool,
             &mut block_pool,
+            &mut const_pool,
         )[0];
 
         let num_src_inputs = input_vars.len();
@@ -64,6 +67,7 @@ impl Transform {
             &mut var_pool,
             &mut def_pool,
             &mut block_pool,
+            &mut const_pool,
         );
 
         // Sanity checks.
@@ -128,6 +132,7 @@ impl Transform {
             var_pool,
             def_pool,
             block_pool,
+            const_pool,
             type_env,
         }
     }
@@ -148,16 +153,17 @@ impl Transform {
 /// pool `var_pool`. If the variable was not present in the symbol table, then add it to the list of
 /// `defined_vars`.
 fn var_index(
-    name: &'static str,
+    name: &str,
     symbol_table: &mut SymbolTable,
     defined_vars: &mut Vec<VarIndex>,
     var_pool: &mut VarPool,
 ) -> VarIndex {
-    match symbol_table.get(name) {
+    let name = name.to_string();
+    match symbol_table.get(&name) {
         Some(&existing_var) => existing_var,
         None => {
             // Materialize the variable.
-            let new_var = var_pool.create(name);
+            let new_var = var_pool.create(name.clone());
             symbol_table.insert(name, new_var);
             defined_vars.push(new_var);
             new_var
@@ -176,7 +182,7 @@ fn rewrite_defined_vars(
 ) -> Vec<VarIndex> {
     let mut new_defined_vars = Vec::new();
     for var in &dummy_def.defined_vars {
-        let own_var = var_index(var.name, symbol_table, defined_vars, var_pool);
+        let own_var = var_index(&var.name, symbol_table, defined_vars, var_pool);
         var_pool.get_mut(own_var).set_def(position, def_index);
         new_defined_vars.push(own_var);
     }
@@ -190,6 +196,7 @@ fn rewrite_expr(
     symbol_table: &mut SymbolTable,
     input_vars: &mut Vec<VarIndex>,
     var_pool: &mut VarPool,
+    const_pool: &mut ConstPool,
 ) -> Apply {
     let (apply_target, dummy_args) = if let DummyExpr::Apply(apply_target, dummy_args) = dummy_expr
     {
@@ -215,7 +222,7 @@ fn rewrite_expr(
     for (i, arg) in dummy_args.into_iter().enumerate() {
         match arg {
             DummyExpr::Var(var) => {
-                let own_var = var_index(var.name, symbol_table, input_vars, var_pool);
+                let own_var = var_index(&var.name, symbol_table, input_vars, var_pool);
                 let var = var_pool.get(own_var);
                 assert!(
                     var.is_input() || var.get_def(position).is_some(),
@@ -227,6 +234,15 @@ fn rewrite_expr(
                 assert!(!apply_target.inst().operands_in[i].is_value());
                 args.push(Expr::Literal(literal));
             }
+            DummyExpr::Constant(constant) => {
+                let const_name = const_pool.insert(constant.0);
+                // Here we abuse var_index by passing an empty, immediately-dropped vector to
+                // `defined_vars`; the reason for this is that unlike the `Var` case above,
+                // constants will create a variable that is not an input variable (it is tracked
+                // instead by ConstPool).
+                let const_var = var_index(&const_name, symbol_table, &mut vec![], var_pool);
+                args.push(Expr::Var(const_var));
+            }
             DummyExpr::Apply(..) => {
                 panic!("Recursive apply is not allowed.");
             }
@@ -248,13 +264,14 @@ fn rewrite_def_list(
     var_pool: &mut VarPool,
     def_pool: &mut DefPool,
     block_pool: &mut BlockPool,
+    const_pool: &mut ConstPool,
 ) -> Vec<DefIndex> {
     let mut new_defs = Vec::new();
     // Register variable names of new blocks first as a block name can be used to jump forward. Thus
     // the name has to be registered first to avoid misinterpreting it as an input-var.
     for dummy_def in dummy_defs.iter() {
         if let DummyExpr::Block(ref var) = dummy_def.expr {
-            var_index(var.name, symbol_table, defined_vars, var_pool);
+            var_index(&var.name, symbol_table, defined_vars, var_pool);
         }
     }
 
@@ -272,7 +289,7 @@ fn rewrite_def_list(
         );
         if let DummyExpr::Block(var) = dummy_def.expr {
             let var_index = *symbol_table
-                .get(var.name)
+                .get(&var.name)
                 .or_else(|| {
                     panic!(
                         "Block {} was not registered during the first visit",
@@ -283,8 +300,14 @@ fn rewrite_def_list(
             var_pool.get_mut(var_index).set_def(position, def_index);
             block_pool.create_block(var_index, def_index);
         } else {
-            let new_apply =
+            let new_apply = rewrite_expr(
-                rewrite_expr(position, dummy_def.expr, symbol_table, input_vars, var_pool);
+                position,
+                dummy_def.expr,
+                symbol_table,
+                input_vars,
+                var_pool,
+                const_pool,
+            );
 
             assert!(
                 def_pool.next_index() == def_index,

cranelift/codegen/meta/src/gen_legalizer.rs

@@ -43,7 +43,7 @@ fn unwrap_inst(
         .args
         .iter()
         .map(|arg| match arg.maybe_var() {
-            Some(var_index) => var_pool.get(var_index).name,
+            Some(var_index) => var_pool.get(var_index).name.as_ref(),
             None => "_",
         })
         .collect::<Vec<_>>()
@@ -114,7 +114,7 @@ fn unwrap_inst(
     assert_eq!(inst.operands_in.len(), apply.args.len());
     for (i, op) in inst.operands_in.iter().enumerate() {
         if op.is_varargs() {
-            let name = var_pool
+            let name = &var_pool
                 .get(apply.args[i].maybe_var().expect("vararg without name"))
                 .name;
 
@@ -283,8 +283,8 @@ fn emit_dst_inst(def: &Def, def_pool: &DefPool, var_pool: &VarPool, fmt: &mut Fo
         let vars = def
             .defined_vars
             .iter()
-            .map(|&var_index| var_pool.get(var_index).name)
+            .map(|&var_index| var_pool.get(var_index).name.as_ref())
-            .collect::<Vec<_>>();
+            .collect::<Vec<&str>>();
         if vars.len() == 1 {
             vars[0].to_string()
         } else {
@@ -401,6 +401,16 @@ fn gen_transform<'a>(
     // Guard the actual expansion by `predicate`.
     fmt.line("if predicate {");
     fmt.indent(|fmt| {
+        // Emit any constants that must be created before use.
+        for (name, value) in transform.const_pool.iter() {
+            fmtln!(
+                fmt,
+                "let {} = pos.func.dfg.constants.insert(vec!{:?});",
+                name,
+                value
+            );
+        }
+
         // If we are adding some blocks, we need to recall the original block, such that we can
         // recompute it.
         if !transform.block_pool.is_empty() {