Update the gensym only when needed during postprocessing

examples/flip-tree/target/output.mz

@@ -2,7 +2,7 @@ main(xs) := f0(xs);
 
 f0(x0) := match x0 {
     Branch(x1, x2) -> +(f1(x1), f1(x2)),
-    Leaf(x3) -> x3
+    Leaf(x1) -> x1
 };
 
 f1(x0) := f0(x0);

examples/imperative-vm/target/output.mz

@@ -22,7 +22,7 @@ f0(x0, x1) := match x1 {
 
 f1(x0) := match x0 {
     Boolean(x1) -> Panic("want `Int(n)`, got `Boolean(b)`"),
-    Int(x2) -> x2
+    Int(x1) -> x1
 };
 
 f10(x0, x1, x2, x3) := match x0 {
@@ -32,8 +32,8 @@ f10(x0, x1, x2, x3) := match x0 {
 
 f11(x0, x1) := match x0 {
     Assign(x2, x3) -> f5(x1, x2, f3(x1, x3)),
-    If(x4, x5, x6) -> f9(f8(f3(x1, x4)), x6, x1, x5),
-    While(x7, x8) -> f10(f8(f3(x1, x7)), x1, x8, While(x7, x8))
+    If(x2, x3, x4) -> f9(f8(f3(x1, x2)), x4, x1, x3),
+    While(x2, x3) -> f10(f8(f3(x1, x2)), x1, x3, While(x2, x3))
 };
 
 f2(x0, x1, x2, x3) := match x3 {
@@ -46,14 +46,12 @@ f2(x0, x1, x2, x3) := match x3 {
 
 f3(x0, x1) := match x1 {
     Boolean(x2) -> Boolean(x2),
-    GreaterThan(x3, x4) ->
-    let x5 := f4(x0, x3); let x6 := f4(x0, x4); Boolean(>(x5, x6)),
-    Int(x7) -> Int(x7),
-    Mul(x8, x9) ->
-    let x10 := f4(x0, x8); let x11 := f4(x0, x9); Int(*(x10, x11)),
-    Sub(x12, x13) ->
-    let x14 := f4(x0, x12); let x15 := f4(x0, x13); Int(-(x14, x15)),
-    Var(x16) -> f0(x16, x0)
+    GreaterThan(x2, x3) ->
+    let x4 := f4(x0, x2); let x5 := f4(x0, x3); Boolean(>(x4, x5)),
+    Int(x2) -> Int(x2),
+    Mul(x2, x3) -> let x4 := f4(x0, x2); let x5 := f4(x0, x3); Int(*(x4, x5)),
+    Sub(x2, x3) -> let x4 := f4(x0, x2); let x5 := f4(x0, x3); Int(-(x4, x5)),
+    Var(x2) -> f0(x2, x0)
 };
 
 f4(x0, x1) := f1(f3(x0, x1));
@@ -69,8 +67,8 @@ f5(x0, x1, x2) := match x0 {
 f6(x0, x1) := match x1 {
     Cons(x2, x3) -> match x2 {
         Assign(x4, x5) -> f7(x3, f5(x0, x4, f3(x0, x5))),
-        If(x6, x7, x8) -> f7(x3, f9(f8(f3(x0, x6)), x8, x0, x7)),
-        While(x9, x10) -> f7(x3, f10(f8(f3(x0, x9)), x0, x10, While(x9, x10)))
+        If(x4, x5, x6) -> f7(x3, f9(f8(f3(x0, x4)), x6, x0, x5)),
+        While(x4, x5) -> f7(x3, f10(f8(f3(x0, x4)), x0, x5, While(x4, x5)))
     },
     Nil() -> x0
 };
@@ -79,7 +77,7 @@ f7(x0, x1) := f6(x1, x0);
 
 f8(x0) := match x0 {
     Boolean(x1) -> x1,
-    Int(x2) -> Panic("want `Boolean(b)`, got `Int(n)`")
+    Int(x1) -> Panic("want `Boolean(b)`, got `Int(n)`")
 };
 
 f9(x0, x1, x2, x3) := match x0 {

examples/productivity-analysis/target/output.mz

@@ -9,12 +9,12 @@ f0(x0) := match x0 {
             T() -> f1(x2)
         }
     },
-    If2(x6, x7, x8) -> match x6 {
-        F() -> f1(x8),
-        T() -> f1(x7),
-        Var(x9, x10) -> match f2(x9, x10) {
-            F() -> f1(x8),
-            T() -> f1(x7)
+    If2(x1, x2, x3) -> match x1 {
+        F() -> f1(x3),
+        T() -> f1(x2),
+        Var(x4, x5) -> match f2(x4, x5) {
+            F() -> f1(x3),
+            T() -> f1(x2)
         }
     }
 };

examples/share-globals/target/output.mz

@@ -2,5 +2,5 @@ main(x, y) := Pair(f0(x), f0(y));
 
 f0(x0) := match x0 {
     A(x1) -> B(x1),
-    C(x2) -> D(x2)
+    C(x1) -> D(x1)
 };

examples/sum-squares-tree/target/output.mz

@@ -2,7 +2,7 @@ main(xs) := f0(xs);
 
 f0(x0) := match x0 {
     Branch(x1, x2) -> +(f1(x1), f1(x2)),
-    Leaf(x3) -> *(x3, x3)
+    Leaf(x1) -> *(x1, x1)
 };
 
 f1(x0) := f0(x0);

examples/zip-map-tree/target/output.mz

@@ -13,11 +13,11 @@ f1(x0, x1, x2, x3) := match x3 {
             Node(x7, x8, x9) -> Node(f0(x4), Pair(Blah(x1), Blah(x5)), f2(x7, x8
             , x9, x6))
         },
-        Node(x10, x11, x12) -> match x2 {
-            Empty() -> Node(f2(x10, x11, x12, x4), Pair(Blah(x1), Blah(x5)), f0(
-            x6)),
-            Node(x13, x14, x15) -> Node(f2(x10, x11, x12, x4), Pair(Blah(x1),
-            Blah(x5)), f2(x13, x14, x15, x6))
+        Node(x7, x8, x9) -> match x2 {
+            Empty() -> Node(f2(x7, x8, x9, x4), Pair(Blah(x1), Blah(x5)), f0(x6)
+            ),
+            Node(x10, x11, x12) -> Node(f2(x7, x8, x9, x4), Pair(Blah(x1), Blah(
+            x5)), f2(x10, x11, x12, x6))
         }
     }
 };

lib/mechanics/postprocessor.ml

@@ -63,7 +63,7 @@ let share_args ~gensym (op, initial_args) =
 let query_env ~env x = Option.value ~default:x (Symbol_map.find_opt x env)
 
 (* TODO: handle other types of restrictions as well. *)
-type restriction = NotEqual of t
+type restriction = NotEqual of Raw_term.t
 
 let handle_term ~(gensym : Gensym.t) ~(env : Renaming.t) t =
     (* We propagate negative information about primitives at residualization-time. During
@@ -76,53 +76,56 @@ let handle_term ~(gensym : Gensym.t) ~(env : Renaming.t) t =
          | _ -> false)
       | _ -> false
     in
+    let hermetic k force =
+        let gensym_backup = Gensym.clone gensym in
+        let result = k force in
+        Gensym.assign ~other:gensym_backup gensym;
+        result
+    in
     let exception Select of Raw_term.t in
     let rec go ~env = function
       | Var x -> Var (query_env ~env x)
       | Const const -> Const const
-      | Call (op, [ Var x; t ]) when op = symbol "=" && is_conflict (x, go_subterm ~env t)
-        -> Call (symbol "F", [])
-      | Call (op, [ t; Var x ]) when op = symbol "=" && is_conflict (x, go_subterm ~env t)
-        -> Call (symbol "F", [])
+      | Call (op, [ Var x; t ])
+        when op = symbol "=" && is_conflict (x, hermetic (go ~env) t) ->
+        Call (symbol "F", [])
+      | Call (op, [ t; Var x ])
+        when op = symbol "=" && is_conflict (x, hermetic (go ~env) t) ->
+        Call (symbol "F", [])
       | Call (op, args) ->
         (* The problem of sharing arguments has optimal substructure: by postprocessing
            arguments before feeding them to [share_args], we make the resulting term share
            all equal arguments which reside on the same syntactical level (a function
            call), throughout all levels in the term. *)
-        share_args ~gensym (op, List.map (go_subterm ~env) args)
+        share_args ~gensym (op, List.map (hermetic (go ~env)) args)
       | Match
           ( (Call (op, ([ Var x; negation ] | [ negation; Var x ])) as t)
           , ([ ((c_f, []), case_f); ((c_t, []), case_t) ] as cases) )
         when op = symbol "=" && c_f = symbol "F" && c_t = symbol "T" ->
         (try
            let t = go_scrutinee ~env ~cases t in
-           let case_f = (c_f, []), go_restrict ~env ~x ~negation case_f in
-           let case_t = (c_t, []), go ~env case_t in
+           let case_f = (c_f, []), hermetic (go_restrict ~env ~x ~negation) case_f in
+           let case_t = (c_t, []), hermetic (go ~env) case_t in
            Raw_term.Match (t, [ case_f; case_t ])
          with
          | Select t -> t)
       | Match (t, cases) ->
         (try
            let t = go_scrutinee ~env ~cases t in
-           let cases = List.map (go_case ~env) cases in
+           let cases = List.map (hermetic (go_case ~env)) cases in
            Raw_term.Match (t, cases)
          with
          | Select t -> t)
       | Let (x, t, u) ->
-        let t = go_subterm ~env t in
+        let t = hermetic (go ~env) t in
         let x' = Gensym.emit gensym in
         let u = go_extend ~env (([ x ], [ x' ]), u) in
         Let (x', t, u)
-    and go_subterm ~env t =
-        let gensym_backup = Gensym.clone gensym in
-        let t = go ~env t in
-        Gensym.assign ~other:gensym_backup gensym;
-        t
     and go_extend ~env ((params, params'), t) =
         let env = Symbol_map.extend2 ~keys:params ~values:params' env in
         go ~env t
     and go_restrict ~env ~x ~negation t =
-        let negation = go_subterm ~env negation in
+        let negation = hermetic (go ~env) negation in
         match negation with
         | Var _ | Const _ ->
           Hashtbl.add restrictions x (NotEqual negation);
@@ -134,15 +137,12 @@ let handle_term ~(gensym : Gensym.t) ~(env : Renaming.t) t =
         let c_params' = Gensym.emit_list ~length_list:c_params gensym in
         (c, c_params'), go_extend ~env ((c_params, c_params'), t)
     and go_scrutinee ~env ~cases t =
-        let gensym_backup = Gensym.clone gensym in
-        let t = go ~env t in
+        let t = hermetic (go ~env) t in
         (match t with
          | Call (c, []) ->
            List.iter
              (function
-               | (c', []), t when c = c' ->
-                 Gensym.assign ~other:gensym_backup gensym;
-                 raise_notrace (Select (go ~env t))
+               | (c', []), t when c = c' -> raise_notrace (Select (go ~env t))
                | _ -> ())
              cases
          | _ -> ());