Zippers struct

src/expansion.rs

@@ -134,7 +134,7 @@ impl std::fmt::Display for ExpandsTo {
 pub fn tracked_expands_to(pattern: &Pattern, hole_zid: ZId, shared: &SharedData) -> ExpandsTo {
     // apply the hole zipper to the original expr being tracked to get the subtree
     // this will expand into, then get the ExpandsTo of that
-    let idx = shared.zip_of_zid[hole_zid].zip(&shared.tracking.as_ref().unwrap().expr);
+    let idx = shared.zippers.zip_of_zid[hole_zid].zip(&shared.tracking.as_ref().unwrap().expr);
     match expands_to_of_node(&shared.tracking.as_ref().unwrap().expr.set[idx]) {
         ExpandsTo(ExpandsToInner::IVar(i, VariableType::Metavar)) => {
             ExpandsTo(ExpandsToInner::IVar(pattern.pattern_args.find_variable(shared, i as usize) as i32, VariableType::Metavar))
@@ -176,7 +176,7 @@ pub fn get_ivars_expansions(original_pattern: &Pattern, arg_of_loc: &FxHashMap<I
 
     if shared.cfg.no_curried_metavars {
         // dont allow any expansions that result in a metavar to the left of an app
-        if shared.zip_of_zid[hole_zid].ends_with_func(){
+        if shared.zippers.zip_of_zid[hole_zid].ends_with_func(){
             return ivars_expansions;
         }
     }

src/pattern_args.rs

@@ -60,7 +60,7 @@ impl PatternArgs {
 
     pub fn use_args(&self, shared: &SharedData, node: &Idx) -> Vec<ZId> {
         self.variables.iter().map(|(zid, _)|
-            shared.arg_of_zid_node[*zid as usize][node].shifted_id
+            shared.zippers.arg_of_zid_node[*zid as usize][node].shifted_id
         ).collect()
     }
 
@@ -77,7 +77,7 @@ impl PatternArgs {
                 continue;
             }
             let zid = self.variables[i].0 as ZId;
-            let shifted_arg = shared.arg_of_zid_node[zid][loc].shifted_id;
+            let shifted_arg = shared.zippers.arg_of_zid_node[zid][loc].shifted_id;
             if !shared.analyzed_ivars[shifted_arg].is_empty() {
                 return true;
             }
@@ -95,8 +95,8 @@ impl PatternArgs {
                     continue;
                 }
                 // if its the same arg in every place, and doesnt have any free vars (ie it's safe to inline)
-                if locs.iter().map(|loc| shared.arg_of_zid_node[argchoice.zid][loc].shifted_id).all_equal()
-                    && shared.analyzed_free_vars[shared.arg_of_zid_node[argchoice.zid][&locs[0]].shifted_id].is_empty()
+                if locs.iter().map(|loc| shared.zippers.arg_of_zid_node[argchoice.zid][loc].shifted_id).all_equal()
+                    && shared.analyzed_free_vars[shared.zippers.arg_of_zid_node[argchoice.zid][&locs[0]].shifted_id].is_empty()
                 {
                     if !shared.cfg.no_stats { shared.stats.lock().deref_mut().useless_abstract_fired += 1; };
                     return true;
@@ -117,13 +117,13 @@ impl PatternArgs {
                 if *type_1 != VariableType::Metavar {
                     continue;
                 }
-                let arg_of_loc_1 = &shared.arg_of_zid_node[*ivar_zid_1 as ZId];
+                let arg_of_loc_1 = &shared.zippers.arg_of_zid_node[*ivar_zid_1 as ZId];
                 // for some reason, the enumerate makes it like 1% faster?????
                 for (_j, (ivar_zid_2, type_2)) in self.variables.iter().enumerate().skip(i+1) {
                     if *type_2 != VariableType::Metavar {
                         continue;
                     }
-                    let arg_of_loc_2 = &shared.arg_of_zid_node[*ivar_zid_2 as ZId];
+                    let arg_of_loc_2 = &shared.zippers.arg_of_zid_node[*ivar_zid_2 as ZId];
                     if locs.iter().all(|loc|
                         arg_of_loc_1[loc].shifted_id == arg_of_loc_2[loc].shifted_id)
                     {
@@ -192,7 +192,7 @@ impl LocationsForReusableArgs<'_> {
 impl PatternArgs {
     pub fn reusable_args_location(&self, shared: &SharedData, ivar: Idx, arg_of_loc: &FxHashMap<Idx, Arg>, match_locations: &mut LocationsForReusableArgs) -> Vec<Idx> {
         let (first_zid_of_var, type_of_var) = self.variables[ivar];
-        let arg_of_loc_ivar = &shared.arg_of_zid_node[first_zid_of_var as ZId];
+        let arg_of_loc_ivar = &shared.zippers.arg_of_zid_node[first_zid_of_var as ZId];
         let relevant_locs = match_locations.relevant_locs(type_of_var, arg_of_loc, &shared.sym_var_info);
         compatible_locations(shared, relevant_locs, arg_of_loc, arg_of_loc_ivar, type_of_var)
     }

src/rewriting.rs

@@ -39,13 +39,13 @@ pub fn rewrite_fast(
            && (!pattern.util_calc.corrected_utils.contains_key(&unshifted_id) // and either we have no conflict (ie corrected_utils doesnt have an entry)
              || pattern.util_calc.corrected_utils[&unshifted_id]) // or we have a conflict but we choose to accept it (which is contextless in this top down approach so its the right move)
         //    && !pattern.pattern.variables.iter().any(|zid| // and there are no negative vars anywhere in the arguments
-        //         shared.egraph[shared.arg_of_zid_node[*zid][&unshifted_id].Idx].data.free_vars.iter().any(|var| *var < 0))
+        //         shared.egraph[shared.zippers.arg_of_zid_node[*zid][&unshifted_id].Idx].data.free_vars.iter().any(|var| *var < 0))
         {
             //  if !shared.cfg.quiet { println!("inv applies at unshifted={} with shift={}", extract(unshifted_id,&shared.egraph), shift) }
             let mut expr = owned_set.add(inv_name.clone());
             // wrap the prim in all the Apps to args
             for zid in pattern.pattern.pattern_args.iterate_one_zid_per_argument() {
-                let arg: &Arg = &shared.arg_of_zid_node[zid][&unshifted_id];
+                let arg: &Arg = &shared.zippers.arg_of_zid_node[zid][&unshifted_id];
 
                 if arg.shift != 0 {
                     shift_rules.push(ShiftRule{depth_cutoff: total_depth, shift: arg.shift});
@@ -66,7 +66,7 @@ pub fn rewrite_fast(
                     // Also note that in the single_hole code --eta-long enforces that match locations never contains anything that starts to the left of a func so
                     // we dont need to worry about the case where the zipper would extend even past the root of the match location
                     // Also note that due to beta normal form, this will be zero and will be a no-op if the arg is a lambda
-                    let arity_of_arg = shared.zip_of_zid[zid].function_arity();
+                    let arity_of_arg = shared.zippers.zip_of_zid[zid].function_arity();
                     if arity_of_arg > 0 {
                         let analyzed_free_vars = &mut AnalyzedExpr::new(FreeVarAnalysis);
 

src/tdfa.rs

@@ -173,7 +173,7 @@ impl TDFAInventionAnnotation {
         let root_sym = global_annotations.symbols[match_location].clone()?;
         let mut ivar_states: Vec<String> = vec![];
         let all_found = pattern.pattern_args.iterate_one_zid_per_argument().all(|ivar_zid| {
-            let Some(node) = shared.arg_of_zid_node[ivar_zid].get(&match_location) else {
+            let Some(node) = shared.zippers.arg_of_zid_node[ivar_zid].get(&match_location) else {
                 return false;
             };
             let Some(ivar_sym) = global_annotations.symbols[node.unshifted_id].clone() else {

src/zipper.rs

@@ -1,5 +1,8 @@
-use lambdas::ExprOwned;
+use lambdas::{AnalyzedExpr, ExprOwned, ExprSet, FreeVarAnalysis, Idx};
 pub use lambdas::{ZNode, ZId, LabelledZId};
+use rustc_hash::FxHashMap;
+
+use crate::{insert_arg_ivars, Arg, Cost, ExpandsTo, ZIdExtension};
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash, Default)]
 pub struct Zipper(Vec<ZNode>);
@@ -47,4 +50,100 @@ impl Zipper {
         expr.immut().zip(&self.0).idx
     }
 
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct Zippers {
+    pub zid_of_zip: FxHashMap<Zipper, ZId>,
+    pub zip_of_zid: Vec<Zipper>,
+    pub arg_of_zid_node: Vec<FxHashMap<Idx,Arg>>,
+}
+
+impl Zippers {
+
+    #[inline(always)]
+    pub fn get_interned_idx(&self, zipper: &Zipper) -> Option<ZId> {
+        self.zid_of_zip.get(zipper).cloned()
+    }
+
+    #[inline(always)]
+    pub fn add_empty(&mut self, empty_zid: ZId) {
+        self.zid_of_zip.insert(Zipper::default(), empty_zid);
+        self.zip_of_zid.push(Zipper::default());
+        self.arg_of_zid_node.push(FxHashMap::default());
+    }
+
+    #[inline(always)]
+    pub fn add_arg(&mut self, zid: ZId, node: Idx, cost: Cost, expands_to: ExpandsTo) { 
+        self.arg_of_zid_node[zid].insert(node,
+            Arg { shifted_id: node, unshifted_id: node, shift: 0, cost, expands_to });
+
+    }
+
+    #[inline(always)]
+    pub fn extend_zipper(&mut self, unextended_zid: ZId, extended_node: Idx, unextended_node: Idx, znode: ZNode) -> usize {
+        let mut zip = self.zip_of_zid[unextended_zid].clone();
+        zip.add_to_front(znode);
+        let zip_of_zid = &mut self.zip_of_zid;
+        let arg_of_zid_node = &mut self.arg_of_zid_node;
+        let zid = self.zid_of_zip.entry(zip.clone()).or_insert_with(|| {
+            let zid = zip_of_zid.len();
+            zip_of_zid.push(zip);
+            arg_of_zid_node.push(FxHashMap::default());
+            zid
+        });
+        // add new zid to this node
+        // give it the same arg
+        let arg = self.arg_of_zid_node[unextended_zid][&unextended_node].clone();
+        self.arg_of_zid_node[*zid].insert(extended_node, arg);
+        *zid
+    }
+
+    #[inline(always)]
+    pub fn handle_shift(&mut self, extended_zid: ZId, unextended_zid: ZId, extended_node: Idx, unextended_node: Idx, analyzed_free_vars: &mut AnalyzedExpr<FreeVarAnalysis>, set: &mut ExprSet) {
+        let zip = &self.zip_of_zid[extended_zid];
+        let mut arg: Arg = self.arg_of_zid_node[unextended_zid][&unextended_node].clone();
+        // shift the arg but keep the unshifted part the same
+        if !analyzed_free_vars.analyze_get(set.get(arg.shifted_id)).is_empty() {
+            // the arg has free vars so we should actually downshift it by 1
+            if analyzed_free_vars[arg.shifted_id].contains(&0) {
+                // furthermore one of those vars is a 0 then it will get shifted to -1, so we handle that slightly specially
+                // by inserting an IVar to indicate this
+
+                // how many lambdas are along this zipper? (including most recent one)
+                let depth_root_to_arg = zip.depth_root_to_arg() as i32;
+
+                // find all pointers to $0 (this is the `init_depth` parameter) and replace then with #(num_lams - 1) that is
+                // point past all lambdas except the newly added one. For example if there were no lambdas other than the
+                // newly added one this would be num_lams=1 so it'd be #0.
+                arg.shifted_id = insert_arg_ivars(&mut set.get_mut(arg.shifted_id), depth_root_to_arg-1, 0, analyzed_free_vars);
+            }
+            arg.shifted_id = set.get_mut(arg.shifted_id).shift(-1, 0, analyzed_free_vars);
+            arg.shift -= 1;
+        }
+        self.arg_of_zid_node[extended_zid].insert(extended_node, arg);
+    }
+
+    #[inline(always)]
+    pub fn compute_extensions(&self) -> Vec<ZIdExtension> {
+        self.zip_of_zid.iter().map(|zip| {
+            let mut zip_body = zip.clone();
+            zip_body.add_to_end(ZNode::Body);
+            let mut zip_arg = zip.clone();
+            zip_arg.add_to_end(ZNode::Arg);
+            let mut zip_func = zip.clone();
+            zip_func.add_to_end(ZNode::Func);
+            ZIdExtension {
+                body: self.zid_of_zip.get(&zip_body).copied(),
+                arg: self.zid_of_zip.get(&zip_arg).copied(),
+                func: self.zid_of_zip.get(&zip_func).copied(),
+            }
+        }).collect()
+    }
+
+    pub fn print_stats(&self) {
+        println!("{} zips", self.zip_of_zid.len());
+        println!("arg_of_zid_node size: {}", self.arg_of_zid_node.len())
+    }
+
 }