Actually use follow types

data/expected_outputs/symbol-reuse-ssm.json

@@ -0,0 +1,122 @@
+{
+  "cmd": "/home/kavi/mit/stitch/target/debug/deps/integration_tests-6cddd9da19e4976e symbols_following",
+  "args": {
+    "iterations": 1,
+    "abstraction_prefix": "fn_",
+    "previous_abstractions": 0,
+    "shuffle": false,
+    "truncate": null,
+    "no_opt": false,
+    "silent": false,
+    "verbose_rewrite": false,
+    "step": {
+      "max_arity": 6,
+      "threads": 1,
+      "no_stats": false,
+      "batch": 1,
+      "dynamic_batch": false,
+      "eta_long": false,
+      "no_curried_metavars": false,
+      "no_curried_bodies": false,
+      "inv_candidates": 1,
+      "hole_choice": "DepthFirst",
+      "cost": {
+        "cost_lam": 1,
+        "cost_app": 1,
+        "cost_var": 100,
+        "cost_ivar": 100,
+        "cost_prim_default": 100,
+        "cost_prim": "{}"
+      },
+      "no_mismatch_check": false,
+      "follow": "(+ 1 #2 #2 #1 #0 2 3 4)",
+      "follow_types": [
+        "Symvar",
+        "Symvar",
+        "Metavar"
+      ],
+      "follow_prune": false,
+      "verbose_worklist": false,
+      "verbose_best": false,
+      "print_stats": 0,
+      "show_rewritten": false,
+      "rewritten_dreamcoder": false,
+      "rewritten_intermediates": false,
+      "inv_arg_cap": false,
+      "allow_single_task": false,
+      "no_opt_single_use": false,
+      "no_opt_upper_bound": false,
+      "no_opt_force_multiuse": false,
+      "no_opt_useless_abstract": false,
+      "no_opt_arity_zero": false,
+      "no_other_util": false,
+      "structure_penalty": 1.0,
+      "rewrite_check": false,
+      "utility_by_rewrite": false,
+      "dreamcoder_comparison": false,
+      "quiet": false,
+      "fused_lambda_tags": {
+        "tags": null
+      },
+      "tdfa": {
+        "tdfa_json_path": null,
+        "tdfa_root": null,
+        "valid_metavars": null,
+        "valid_roots": null,
+        "tdfa_non_eta_long_states": null,
+        "tdfa_split": null
+      },
+      "symvar": {
+        "symvar_prefix": "&"
+      }
+    }
+  },
+  "original_cost": 2724,
+  "final_cost": 1209,
+  "compression_ratio": 2.2531017369727047,
+  "num_abstractions": 1,
+  "original": [
+    "(+ 1 &x &x &y &z 2 3 4)",
+    "(+ 1 &y &y &x &z 2 3 4)",
+    "(+ 1 &x &x &x &z 2 3 4)"
+  ],
+  "rewritten": [
+    "(fn_0 &z &y &x)",
+    "(fn_0 &z &x &y)",
+    "(fn_0 &z &x &x)"
+  ],
+  "rewritten_dreamcoder": null,
+  "abstractions": [
+    {
+      "body": "(+ 1 #2 #2 #1 #0 2 3 4)",
+      "dreamcoder": "#(lambda (lambda (lambda (+ 1 $0 $0 $1 $2 2 3 4))))",
+      "arity": 3,
+      "name": "fn_0",
+      "utility": 1007,
+      "final_cost": 1209,
+      "compression_ratio": 2.2531017369727047,
+      "cumulative_compression_ratio": 2.2531017369727047,
+      "num_uses": 3,
+      "rewritten": null,
+      "rewritten_dreamcoder": null,
+      "uses": [
+        {
+          "fn_0 &z &x &x": "(+ 1 &x &x &x &z 2 3 4)"
+        },
+        {
+          "fn_0 &z &y &x": "(+ 1 &x &x &y &z 2 3 4)"
+        },
+        {
+          "fn_0 &z &x &y": "(+ 1 &y &y &x &z 2 3 4)"
+        }
+      ],
+      "dc_comparison_millis": null,
+      "tdfa_annotation": null,
+      "variable_types": [
+        "S",
+        "M",
+        "M"
+      ]
+    }
+  ]
+}

data/expected_outputs/symbol-reuse-sss.json

@@ -0,0 +1,122 @@
+{
+  "cmd": "/home/kavi/mit/stitch/target/debug/deps/integration_tests-6cddd9da19e4976e symbols_following",
+  "args": {
+    "iterations": 1,
+    "abstraction_prefix": "fn_",
+    "previous_abstractions": 0,
+    "shuffle": false,
+    "truncate": null,
+    "no_opt": false,
+    "silent": false,
+    "verbose_rewrite": false,
+    "step": {
+      "max_arity": 6,
+      "threads": 1,
+      "no_stats": false,
+      "batch": 1,
+      "dynamic_batch": false,
+      "eta_long": false,
+      "no_curried_metavars": false,
+      "no_curried_bodies": false,
+      "inv_candidates": 1,
+      "hole_choice": "DepthFirst",
+      "cost": {
+        "cost_lam": 1,
+        "cost_app": 1,
+        "cost_var": 100,
+        "cost_ivar": 100,
+        "cost_prim_default": 100,
+        "cost_prim": "{}"
+      },
+      "no_mismatch_check": false,
+      "follow": "(+ 1 #2 #2 #1 #0 2 3 4)",
+      "follow_types": [
+        "Symvar",
+        "Symvar",
+        "Symvar"
+      ],
+      "follow_prune": false,
+      "verbose_worklist": false,
+      "verbose_best": false,
+      "print_stats": 0,
+      "show_rewritten": false,
+      "rewritten_dreamcoder": false,
+      "rewritten_intermediates": false,
+      "inv_arg_cap": false,
+      "allow_single_task": false,
+      "no_opt_single_use": false,
+      "no_opt_upper_bound": false,
+      "no_opt_force_multiuse": false,
+      "no_opt_useless_abstract": false,
+      "no_opt_arity_zero": false,
+      "no_other_util": false,
+      "structure_penalty": 1.0,
+      "rewrite_check": false,
+      "utility_by_rewrite": false,
+      "dreamcoder_comparison": false,
+      "quiet": false,
+      "fused_lambda_tags": {
+        "tags": null
+      },
+      "tdfa": {
+        "tdfa_json_path": null,
+        "tdfa_root": null,
+        "valid_metavars": null,
+        "valid_roots": null,
+        "tdfa_non_eta_long_states": null,
+        "tdfa_split": null
+      },
+      "symvar": {
+        "symvar_prefix": "&"
+      }
+    }
+  },
+  "original_cost": 2724,
+  "final_cost": 1209,
+  "compression_ratio": 2.2531017369727047,
+  "num_abstractions": 1,
+  "original": [
+    "(+ 1 &x &x &y &z 2 3 4)",
+    "(+ 1 &y &y &x &z 2 3 4)",
+    "(+ 1 &x &x &x &z 2 3 4)"
+  ],
+  "rewritten": [
+    "(fn_0 &z &y &x)",
+    "(fn_0 &z &x &y)",
+    "(fn_0 &z &x &x)"
+  ],
+  "rewritten_dreamcoder": null,
+  "abstractions": [
+    {
+      "body": "(+ 1 #2 #2 #1 #0 2 3 4)",
+      "dreamcoder": "#(lambda (lambda (lambda (+ 1 $0 $0 $1 $2 2 3 4))))",
+      "arity": 3,
+      "name": "fn_0",
+      "utility": 1007,
+      "final_cost": 1209,
+      "compression_ratio": 2.2531017369727047,
+      "cumulative_compression_ratio": 2.2531017369727047,
+      "num_uses": 3,
+      "rewritten": null,
+      "rewritten_dreamcoder": null,
+      "uses": [
+        {
+          "fn_0 &z &x &x": "(+ 1 &x &x &x &z 2 3 4)"
+        },
+        {
+          "fn_0 &z &y &x": "(+ 1 &x &x &y &z 2 3 4)"
+        },
+        {
+          "fn_0 &z &x &y": "(+ 1 &y &y &x &z 2 3 4)"
+        }
+      ],
+      "dc_comparison_millis": null,
+      "tdfa_annotation": null,
+      "variable_types": [
+        "S",
+        "M",
+        "M"
+      ]
+    }
+  ]
+}

src/compression.rs

@@ -656,6 +656,7 @@ pub struct SharedData {
     pub multistep_cfg: MultistepCompressionConfig,
     pub tracking: Option<Tracking>,
     pub fused_lambda_tags: Option<FxHashSet<Tag>>,
+    pub follow: Option<Invention>,
 }
 
 pub fn invalid_metavar_location(shared : &SharedData, node: Idx) -> bool {
@@ -2137,6 +2138,7 @@ pub fn construct_shared(
         multistep_cfg: multistep_cfg.clone(),
         tracking,
         fused_lambda_tags: fused_copy,
+        follow: follow.clone(),
     });
 
     if !shared.cfg.quiet { println!("built SharedData: {:?}ms", tstart.elapsed().as_millis()) }

src/expansion.rs

@@ -134,10 +134,12 @@ 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.tracking.as_ref().unwrap().expr.immut().zip(&shared.zip_of_zid[hole_zid]).idx;
+    let zip = &shared.zip_of_zid[hole_zid];
+    let idx = shared.tracking.as_ref().unwrap().expr.immut().zip(zip).idx;
     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))
+        ExpandsTo(ExpandsToInner::IVar(i, _)) => {
+            let vt = shared.follow.as_ref().unwrap().variable_types[i as usize];
+            ExpandsTo(ExpandsToInner::IVar(pattern.pattern_args.find_variable(shared, i as usize) as i32, vt))
         }
         e => e
     }
@@ -151,7 +153,7 @@ pub fn expands_to_of_node(node: &Node) -> ExpandsTo {
             Node::Prim(p) => ExpandsToInner::Prim(p.clone()),
             Node::Lam(_, tag) => ExpandsToInner::Lam(*tag),
             Node::App(_,_) => ExpandsToInner::App,
-            Node::IVar(i) => ExpandsToInner::IVar(*i, VariableType::Metavar),
+            Node::IVar(i) => ExpandsToInner::IVar(*i, /*placeholder, will be determined by caller*/ VariableType::Metavar),
         }
     )
 }

tests/integration_tests.rs

@@ -142,12 +142,19 @@ fn python_symbols_regression() {
 
 #[test]
 fn symbols_basic() {
-    compare_out_jsons_testing("data/python/symbols-alignment.json", "data/expected_outputs/symbols-alignment.json", "-i2 -a3 --symvar-prefix & ", InputFormat::ProgramsList);
-    compare_out_jsons_testing("data/python/symbols-cannot-be-literal.json", "data/expected_outputs/symbols-cannot-be-literal.json", "-i2 -a3 --symvar-prefix & ", InputFormat::ProgramsList);
-    compare_out_jsons_testing("data/python/symbols-cannot-be-literal-0-arity.json", "data/expected_outputs/symbols-cannot-be-literal-0-arity.json", "-i2 -a3 --symvar-prefix & ", InputFormat::ProgramsList);
-    compare_out_jsons_testing("data/python/symbol-reuse.json", "data/expected_outputs/symbol-reuse.json", "-i1 -a0 --symvar-prefix & ", InputFormat::ProgramsList);
-    compare_out_jsons_testing("data/python/symbol-reuse-dfa.json", "data/expected_outputs/symbol-reuse-dfa.json", &("-i1 -a0 ".to_owned() + &python_args()), InputFormat::ProgramsList);
-    compare_out_jsons_testing("data/python/pick-up-on-abstractions-0-arity.json", "data/expected_outputs/pick-up-on-abstractions-0-arity.json", &("-i1 -a0 ".to_owned() + &python_args()), InputFormat::ProgramsList);
+    compare_out_jsons_testing("data/python/symbols-alignment.json", "data/expected_outputs/symbols-alignment.json", "-i2 -a3  --rewrite-check --symvar-prefix & ", InputFormat::ProgramsList);
+    compare_out_jsons_testing("data/python/symbols-cannot-be-literal.json", "data/expected_outputs/symbols-cannot-be-literal.json", "-i2 -a3  --rewrite-check --symvar-prefix & ", InputFormat::ProgramsList);
+    compare_out_jsons_testing("data/python/symbols-cannot-be-literal-0-arity.json", "data/expected_outputs/symbols-cannot-be-literal-0-arity.json", "-i2 -a3  --rewrite-check --symvar-prefix & ", InputFormat::ProgramsList);
+    compare_out_jsons_testing("data/python/symbol-reuse.json", "data/expected_outputs/symbol-reuse.json", "-i1 -a0 --rewrite-check --symvar-prefix & ", InputFormat::ProgramsList);
+    compare_out_jsons_testing("data/python/symbol-reuse-dfa.json", "data/expected_outputs/symbol-reuse-dfa.json", &("-i1 -a0 --rewrite-check ".to_owned() + &python_args()), InputFormat::ProgramsList);
+    compare_out_jsons_testing("data/python/pick-up-on-abstractions-0-arity.json", "data/expected_outputs/pick-up-on-abstractions-0-arity.json", &("-i1 -a0 --rewrite-check ".to_owned() + &python_args()), InputFormat::ProgramsList);
+}
+
+#[test]
+fn symbols_following() {
+    compare_out_jsons_testing("data/python/symbol-reuse.json", "data/expected_outputs/symbol-reuse.json", "-i1 -a0 --symvar-prefix & --rewrite-check ", InputFormat::ProgramsList);
+    compare_out_jsons_testing("data/python/symbol-reuse.json", "data/expected_outputs/symbol-reuse-sss.json", "-i1 -a6 --symvar-prefix & --follow '(+ 1 #2 #2 #1 #0 2 3 4)' --follow-types 'S S S' --rewrite-check ", InputFormat::ProgramsList);
+    compare_out_jsons_testing("data/python/symbol-reuse.json", "data/expected_outputs/symbol-reuse-ssm.json", "-i1 -a6 --symvar-prefix & --follow '(+ 1 #2 #2 #1 #0 2 3 4)' --follow-types 'S S M' --rewrite-check ", InputFormat::ProgramsList);
 }
 
 #[test]