Add (dynamic) constants and types to editor

hercules_opt/src/dce.rs

@@ -25,7 +25,7 @@ pub fn dce(editor: &mut FunctionEditor) {
 
         // If a node on the worklist has 0 users, delete it. Add its uses onto
         // the worklist.
-        if editor.users(work).len() == 0 {
+        if editor.get_users(work).len() == 0 {
             let uses = get_uses(&editor.func().nodes[work.idx()]);
             let success = editor.edit(|edit| edit.delete_node(work));
             if success {

hercules_opt/src/editor.rs

@@ -3,9 +3,11 @@ extern crate either;
 extern crate hercules_ir;
 extern crate itertools;
 
+use std::cell::{Ref, RefCell};
 use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
 use std::iter::FromIterator;
 use std::mem::take;
+use std::ops::Deref;
 
 use self::bitvec::prelude::*;
 use self::either::Either;
@@ -26,12 +28,19 @@ pub type Edit = (HashSet<NodeID>, HashSet<NodeID>);
 /*
  * Helper object for editing Hercules functions in a trackable manner. Edits are
  * recorded in order to repair partitions and debug info.
+ * Edits must be made atomically, that is, only one `.edit` may be called at a time
+ * across all editors.
  */
 #[derive(Debug)]
 pub struct FunctionEditor<'a> {
     // Wraps a mutable reference to a function. Doesn't provide access to this
     // reference directly, so that we can monitor edits.
     function: &'a mut Function,
+    // Keep a RefCell to (dynamic) constants and types to allow function changes
+    // to update these
+    constants: &'a RefCell<Vec<Constant>>,
+    dynamic_constants: &'a RefCell<Vec<DynamicConstant>>,
+    types: &'a RefCell<Vec<Type>>,
     // Most optimizations need def use info, so provide an iteratively updated
     // mutable version that's automatically updated based on recorded edits.
     mut_def_use: Vec<HashSet<NodeID>>,
@@ -63,17 +72,28 @@ pub struct FunctionEdit<'a: 'b, 'b> {
     // Reference the active function editor.
     editor: &'b mut FunctionEditor<'a>,
     // Keep track of deleted node IDs.
-    deleted: HashSet<NodeID>,
+    deleted_nodeids: HashSet<NodeID>,
     // Keep track of added node IDs.
-    added: HashSet<NodeID>,
+    added_nodeids: HashSet<NodeID>,
     // Keep track of added and use updated nodes.
-    added_and_updated: BTreeMap<NodeID, Node>,
+    added_and_updated_nodes: BTreeMap<NodeID, Node>,
+    // Keep track of added (dynamic) constants and types
+    added_constants: Vec<Constant>,
+    added_dynamic_constants: Vec<DynamicConstant>,
+    added_types: Vec<Type>,
     // Compute a def-use map entries iteratively.
     updated_def_use: BTreeMap<NodeID, HashSet<NodeID>>,
+    updated_return_type: Option<TypeID>,
 }
 
 impl<'a: 'b, 'b> FunctionEditor<'a> {
-    pub fn new(function: &'a mut Function, def_use: &ImmutableDefUseMap) -> Self {
+    pub fn new(
+        function: &'a mut Function,
+        constants: &'a RefCell<Vec<Constant>>,
+        dynamic_constants: &'a RefCell<Vec<DynamicConstant>>,
+        types: &'a RefCell<Vec<Type>>,
+        def_use: &ImmutableDefUseMap,
+    ) -> Self {
         let mut_def_use = (0..function.nodes.len())
             .map(|idx| {
                 def_use
@@ -87,6 +107,9 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
 
         FunctionEditor {
             function,
+            constants,
+            dynamic_constants,
+            types,
             mut_def_use,
             edits: vec![],
             mutable_nodes,
@@ -100,10 +123,14 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
         // Create the edit helper struct and perform the edit using it.
         let edit_obj = FunctionEdit {
             editor: self,
-            deleted: HashSet::new(),
-            added: HashSet::new(),
-            added_and_updated: BTreeMap::new(),
+            deleted_nodeids: HashSet::new(),
+            added_nodeids: HashSet::new(),
+            added_constants: Vec::new().into(),
+            added_dynamic_constants: Vec::new().into(),
+            added_types: Vec::new().into(),
+            added_and_updated_nodes: BTreeMap::new(),
             updated_def_use: BTreeMap::new(),
+            updated_return_type: None,
         };
 
         if let Ok(populated_edit) = edit(edit_obj) {
@@ -111,10 +138,14 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
             // without modifying immutable nodes.
             let FunctionEdit {
                 editor,
-                deleted,
-                added,
-                added_and_updated,
+                deleted_nodeids,
+                added_nodeids,
+                added_constants,
+                added_dynamic_constants,
+                added_types,
+                added_and_updated_nodes: added_and_updated,
                 updated_def_use,
+                updated_return_type,
             } = populated_edit;
             // Step 1: update the mutable def use map.
             for (u, new_users) in updated_def_use {
@@ -148,14 +179,14 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
             // Step 3: delete nodes. This is done using "gravestones", where a
             // node other than node ID 0 being a start node is considered a
             // gravestone.
-            for id in deleted.iter() {
+            for id in deleted_nodeids.iter() {
                 // Check that there are no users of deleted nodes.
                 assert!(editor.mut_def_use[id.idx()].is_empty());
                 editor.function.nodes[id.idx()] = Node::Start;
             }
 
             // Step 4: add a single edit to the edit list.
-            editor.edits.push((deleted, added));
+            editor.edits.push((deleted_nodeids, added_nodeids));
 
             // Step 5: update the length of mutable_nodes. All added nodes are
             // mutable.
@@ -163,6 +194,20 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
                 .mutable_nodes
                 .resize(editor.function.nodes.len(), true);
 
+            // Step 6: update types and constants
+            let mut editor_constants = editor.constants.borrow_mut();
+            let mut editor_dynamic_constants = editor.dynamic_constants.borrow_mut();
+            let mut editor_types = editor.types.borrow_mut();
+
+            editor_constants.extend(added_constants);
+            editor_dynamic_constants.extend(added_dynamic_constants);
+            editor_types.extend(added_types);
+
+            // Step 7: update return type if necessary
+            if let Some(return_type) = updated_return_type {
+                editor.function.return_type = return_type;
+            }
+
             true
         } else {
             false
@@ -173,7 +218,7 @@ impl<'a: 'b, 'b> FunctionEditor<'a> {
         &self.function
     }
 
-    pub fn users(&self, id: NodeID) -> impl ExactSizeIterator<Item = NodeID> + '_ {
+    pub fn get_users(&self, id: NodeID) -> impl ExactSizeIterator<Item = NodeID> + '_ {
         self.mut_def_use[id.idx()].iter().map(|x| *x)
     }
 
@@ -200,7 +245,7 @@ impl<'a, 'b> FunctionEdit<'a, 'b> {
     }
 
     pub fn add_node(&mut self, node: Node) -> NodeID {
-        let id = NodeID::new(self.editor.function.nodes.len() + self.added.len());
+        let id = NodeID::new(self.editor.function.nodes.len() + self.added_nodeids.len());
         // Added nodes need to have an entry in the def-use map.
         self.updated_def_use.insert(id, HashSet::new());
         // Added nodes use other nodes, and we need to update their def-use
@@ -210,8 +255,8 @@ impl<'a, 'b> FunctionEdit<'a, 'b> {
             self.updated_def_use.get_mut(u).unwrap().insert(id);
         }
         // Add the node.
-        self.added_and_updated.insert(id, node);
-        self.added.insert(id);
+        self.added_and_updated_nodes.insert(id, node);
+        self.added_nodeids.insert(id);
         id
     }
 
@@ -220,7 +265,7 @@ impl<'a, 'b> FunctionEdit<'a, 'b> {
         // immutable node, as it means the whole edit should be aborted.
         if self.editor.mutable_nodes[id.idx()] {
             assert!(
-                !self.added.contains(&id),
+                !self.added_nodeids.contains(&id),
                 "PANIC: Please don't delete a node that was added in the same edit."
             );
             // Deleted nodes use other nodes, and we need to update their def-
@@ -232,7 +277,7 @@ impl<'a, 'b> FunctionEdit<'a, 'b> {
                 self.ensure_updated_def_use_entry(u);
                 self.updated_def_use.get_mut(&u).unwrap().remove(&id);
             }
-            self.deleted.insert(id);
+            self.deleted_nodeids.insert(id);
             Ok(self)
         } else {
             Err(self)
@@ -248,14 +293,14 @@ impl<'a, 'b> FunctionEdit<'a, 'b> {
             self.ensure_updated_def_use_entry(old);
             for user_id in self.updated_def_use[&old].iter() {
                 // Replace uses of old with new.
-                let mut updated_user = self.node(*user_id).clone();
+                let mut updated_user = self.get_node(*user_id).clone();
                 for u in get_uses_mut(&mut updated_user).as_mut() {
                     if **u == old {
                         **u = new;
                     }
                 }
                 // Add the updated user to added_and_updated.
-                self.added_and_updated.insert(*user_id, updated_user);
+                self.added_and_updated_nodes.insert(*user_id, updated_user);
             }
 
             // All of the users of the old node become users of the new node, so
@@ -273,9 +318,9 @@ impl<'a, 'b> FunctionEdit<'a, 'b> {
         }
     }
 
-    pub fn node(&self, id: NodeID) -> &Node {
-        assert!(!self.deleted.contains(&id));
-        if let Some(node) = self.added_and_updated.get(&id) {
+    pub fn get_node(&self, id: NodeID) -> &Node {
+        assert!(!self.deleted_nodeids.contains(&id));
+        if let Some(node) = self.added_and_updated_nodes.get(&id) {
             // Refer to added or updated node. This node is guaranteed to be
             // updated with uses after replace_all_uses is called.
             node
@@ -285,8 +330,8 @@ impl<'a, 'b> FunctionEdit<'a, 'b> {
         }
     }
 
-    pub fn users(&self, id: NodeID) -> impl Iterator<Item = NodeID> + '_ {
-        assert!(!self.deleted.contains(&id));
+    pub fn get_users(&self, id: NodeID) -> impl Iterator<Item = NodeID> + '_ {
+        assert!(!self.deleted_nodeids.contains(&id));
         if let Some(users) = self.updated_def_use.get(&id) {
             // Refer to the updated users set.
             Either::Left(users.iter().map(|x| *x))
@@ -295,6 +340,107 @@ impl<'a, 'b> FunctionEdit<'a, 'b> {
             Either::Right(self.editor.mut_def_use[id.idx()].iter().map(|x| *x))
         }
     }
+
+    pub fn add_type(&mut self, ty: Type) -> TypeID {
+        let pos = self
+            .editor
+            .types
+            .borrow()
+            .iter()
+            .chain(self.added_types.iter())
+            .position(|t| *t == ty);
+        if let Some(idx) = pos {
+            TypeID::new(idx)
+        } else {
+            let id = TypeID::new(self.editor.types.borrow().len() + self.added_types.len());
+            self.added_types.push(ty);
+            id
+        }
+    }
+
+    pub fn get_type(&self, id: TypeID) -> impl Deref + '_ {
+        if id.idx() < self.editor.types.borrow().len() {
+            Either::Left(Ref::map(self.editor.types.borrow(), |types| {
+                &types[id.idx()]
+            }))
+        } else {
+            Either::Right(
+                self.added_types
+                    .get(id.idx() - self.editor.types.borrow().len())
+                    .unwrap(),
+            )
+        }
+    }
+
+    pub fn add_constant(&mut self, constant: Constant) -> ConstantID {
+        let pos = self
+            .editor
+            .constants
+            .borrow()
+            .iter()
+            .chain(self.added_constants.iter())
+            .position(|c| *c == constant);
+        if let Some(idx) = pos {
+            ConstantID::new(idx)
+        } else {
+            let id =
+                ConstantID::new(self.editor.constants.borrow().len() + self.added_constants.len());
+            self.added_constants.push(constant);
+            id
+        }
+    }
+
+    pub fn get_constant(&self, id: ConstantID) -> impl Deref + '_ {
+        if id.idx() < self.editor.constants.borrow().len() {
+            Either::Left(Ref::map(self.editor.constants.borrow(), |constants| {
+                &constants[id.idx()]
+            }))
+        } else {
+            Either::Right(
+                self.added_constants
+                    .get(id.idx() - self.editor.constants.borrow().len())
+                    .unwrap(),
+            )
+        }
+    }
+
+    pub fn add_dynamic_constant(&mut self, dynamic_constant: DynamicConstant) -> DynamicConstantID {
+        let pos = self
+            .editor
+            .dynamic_constants
+            .borrow()
+            .iter()
+            .chain(self.added_dynamic_constants.iter())
+            .position(|c| *c == dynamic_constant);
+        if let Some(idx) = pos {
+            DynamicConstantID::new(idx)
+        } else {
+            let id = DynamicConstantID::new(
+                self.editor.dynamic_constants.borrow().len() + self.added_dynamic_constants.len(),
+            );
+            self.added_dynamic_constants.push(dynamic_constant);
+            id
+        }
+    }
+
+    pub fn get_dynamic_constant(&self, id: DynamicConstantID) -> impl Deref + '_ {
+        if id.idx() < self.editor.dynamic_constants.borrow().len() {
+            Either::Left(Ref::map(
+                self.editor.dynamic_constants.borrow(),
+                |dynamic_constants| &dynamic_constants[id.idx()],
+            ))
+        } else {
+            Either::Right(
+                self.added_dynamic_constants
+                    .get(id.idx() - self.editor.dynamic_constants.borrow().len())
+                    .unwrap(),
+            )
+        }
+    }
+
+    pub fn set_return_type(&mut self, ty: TypeID) {
+        self.updated_return_type = Some(ty);
+    }
 }
 
 /*
@@ -606,8 +752,17 @@ fn func(x: i32) -> i32
             .next()
             .unwrap();
 
+        let constants_ref = RefCell::new(src_module.constants);
+        let dynamic_constants_ref = RefCell::new(src_module.dynamic_constants);
+        let types_ref = RefCell::new(src_module.types);
         // Edit the function by replacing the add with a multiply.
-        let mut editor = FunctionEditor::new(func, &def_use(func));
+        let mut editor = FunctionEditor::new(
+            func,
+            &constants_ref,
+            &dynamic_constants_ref,
+            &types_ref,
+            &def_use(func),
+        );
         let success = editor.edit(|mut edit| {
             let mul = edit.add_node(Node::Binary {
                 op: BinaryOperator::Mul,

hercules_opt/src/gvn.rs

@@ -28,7 +28,7 @@ pub fn gvn(editor: &mut FunctionEditor, constants: &Vec<Constant>) {
             }
 
             // Record the users of `work` before making any edits.
-            let work_users: Vec<NodeID> = editor.users(work).collect();
+            let work_users: Vec<NodeID> = editor.get_users(work).collect();
 
             // At this point, we know the number of the node IDed `work` is
             // `number`. We want to replace `work` with `number`, which means

hercules_opt/src/pass.rs

@@ -4,6 +4,7 @@ extern crate postcard;
 extern crate serde;
 extern crate take_mut;
 
+use std::cell::RefCell;
 use std::collections::HashMap;
 use std::env::temp_dir;
 use std::fs::File;
@@ -330,10 +331,24 @@ impl PassManager {
                     self.make_def_uses();
                     let def_uses = self.def_uses.as_ref().unwrap();
                     for idx in 0..self.module.functions.len() {
-                        let mut editor =
-                            FunctionEditor::new(&mut self.module.functions[idx], &def_uses[idx]);
+                        let constants_ref =
+                            RefCell::new(std::mem::take(&mut self.module.constants));
+                        let dynamic_constants_ref =
+                            RefCell::new(std::mem::take(&mut self.module.dynamic_constants));
+                        let types_ref = RefCell::new(std::mem::take(&mut self.module.types));
+                        let mut editor = FunctionEditor::new(
+                            &mut self.module.functions[idx],
+                            &constants_ref,
+                            &dynamic_constants_ref,
+                            &types_ref,
+                            &def_uses[idx],
+                        );
                         dce(&mut editor);
 
+                        self.module.constants = constants_ref.take();
+                        self.module.dynamic_constants = dynamic_constants_ref.take();
+                        self.module.types = types_ref.take();
+
                         let edits = &editor.edits();
                         if let Some(plans) = self.plans.as_mut() {
                             repair_plan(&mut plans[idx], &self.module.functions[idx], edits);
@@ -365,10 +380,24 @@ impl PassManager {
                     self.make_def_uses();
                     let def_uses = self.def_uses.as_ref().unwrap();
                     for idx in 0..self.module.functions.len() {
-                        let mut editor =
-                            FunctionEditor::new(&mut self.module.functions[idx], &def_uses[idx]);
+                        let constants_ref =
+                            RefCell::new(std::mem::take(&mut self.module.constants));
+                        let dynamic_constants_ref =
+                            RefCell::new(std::mem::take(&mut self.module.dynamic_constants));
+                        let types_ref = RefCell::new(std::mem::take(&mut self.module.types));
+                        let mut editor = FunctionEditor::new(
+                            &mut self.module.functions[idx],
+                            &constants_ref,
+                            &dynamic_constants_ref,
+                            &types_ref,
+                            &def_uses[idx],
+                        );
                         gvn(&mut editor, &self.module.constants);
 
+                        self.module.constants = constants_ref.take();
+                        self.module.dynamic_constants = dynamic_constants_ref.take();
+                        self.module.types = types_ref.take();
+
                         let edits = &editor.edits();
                         if let Some(plans) = self.plans.as_mut() {
                             repair_plan(&mut plans[idx], &self.module.functions[idx], edits);