Merging individual methods to one shared method

src/arr/compiler/anf.arr

@@ -223,7 +223,7 @@ fun anf(e :: A.Expr, k :: ANFCont) -> N.AExpr:
       end
       anf(A.s-let-expr(l, let-binds, A.s-block(l, assigns + [list: body]), true), k)
 
-    | s-data-expr(l, data-name, data-name-t, params, mixins, variants, shared, _check-loc, _check) =>
+    | s-data-expr(l, data-name, data-name-type, data-name-ann, params, mixins, variants, shared, _check-loc, _check) =>
       fun anf-member(member :: A.VariantMember):
         cases(A.VariantMember) member:
           | s-variant-member(l2, typ, b) =>
@@ -272,7 +272,7 @@ fun anf(e :: A.Expr, k :: ANFCont) -> N.AExpr:
               N.a-field(f.l, f.name, t)
             end
           anf-variants(variants, lam(new-variants):
-              k(N.a-data-expr(l, data-name, data-name-t, new-variants, new-shared))
+              k(N.a-data-expr(l, data-name, data-name-type, new-variants, new-shared))
             end)
         end)
 

src/arr/compiler/ast-util.arr

@@ -246,12 +246,12 @@ fun default-env-map-visitor<a, c>(
     method s-singleton-cases-branch(self, l, pat-loc, name, body):
       A.s-singleton-cases-branch(l, pat-loc, name, body.visit(self))
     end,
-    method s-data-expr(self, l, name, namet, params, mixins, variants, shared-members, _check-loc, _check):
+    method s-data-expr(self, l, name, name-type, name-ann, params, mixins, variants, shared-members, _check-loc, _check):
       new-type-env = for lists.fold(acc from self.type-env, param from params):
         bind-handlers.s-param-bind(l, param, acc)
       end
       with-params = self.{type-env: new-type-env}
-      A.s-data-expr(l, name, namet.visit(with-params), params,
+      A.s-data-expr(l, name, name-type.visit(with-params), name-ann.visit(with-params), params,
         mixins.map(_.visit(with-params)), variants.map(_.visit(with-params)),
         shared-members.map(_.visit(with-params)), _check-loc, with-params.option(_check))
     end,
@@ -366,12 +366,12 @@ fun default-env-iter-visitor<a, c>(
       and body.visit(self.{env: args-env})
     end,
     # s-singleton-cases-branch introduces no new bindings, so default visitor is fine
-    method s-data-expr(self, l, name, namet, params, mixins, variants, shared-members, _check-loc, _check):
+    method s-data-expr(self, l, name, name-type, name-ann, params, mixins, variants, shared-members, _check-loc, _check):
       new-type-env = for lists.fold(acc from self.type-env, param from params):
         bind-handlers.s-param-bind(l, param, acc)
       end
       with-params = self.{type-env: new-type-env}
-      namet.visit(with-params)
+      name-type.visit(with-params) and name-ann.visit(with-params)
       and lists.all(_.visit(with-params), mixins)
       and lists.all(_.visit(with-params), variants)
       and lists.all(_.visit(with-params), shared-members)
@@ -1052,7 +1052,7 @@ fun get-named-provides(resolved :: CS.NameResolution, uri :: URI, compile-env ::
   end
   fun data-expr-to-datatype(exp :: A.Expr % (is-s-data-expr)) -> T.DataType:
     cases(A.Expr) exp:
-      | s-data-expr(l, name, _, params, _, variants, shared-members, _, _) =>
+      | s-data-expr(l, name, _, _, params, _, variants, shared-members, _, _) =>
 
         tvars = for map(tvar from params):
           T.t-var(tvar, l, false)

src/arr/compiler/compile-lib.arr

@@ -404,8 +404,12 @@ fun compile-module(locator :: Locator, provide-map :: SD.StringDict<URI>, module
               | ok(_) =>
                 var tc-ast = type-checked.code
                 type-checked := nothing
-                var dp-ast = DP.desugar-post-tc(tc-ast, env)
+                var merged = DP.merge-methods(tc-ast)
                 tc-ast := nothing
+                add-phase("Merged methods", merged.ast)
+                var dp-ast = DP.desugar-post-tc(merged.ast, env)
+                named-result.bindings.merge-now(merged.new-binds)
+                merged := nothing
                 var cleaned = dp-ast
                 dp-ast := nothing
                 cleaned := cleaned.visit(AU.letrec-visitor)

src/arr/compiler/desugar-post-tc.arr

@@ -3,11 +3,15 @@
 provide *
 provide-types *
 import ast as A
+import string-dict as SD
+import file("list-aux.arr") as LA
 import file("desugar.arr") as D
 import file("compile-structs.arr") as C
 
-mk-id = D.mk-id
+names = A.global-names
+
 no-branches-exn = D.no-branches-exn
+is-s-method = A.is-s-method
 flat-prim-app = A.prim-app-info-c(false)
 
 fun no-cases-exn(l, val):
@@ -38,6 +42,176 @@ desugar-visitor = A.default-map-visitor.{
   end
 }
 
+fun no-method-exn(l, obj, name):
+  A.s-prim-app(l, "throwFieldNotFound", [list: A.s-srcloc(l, l), obj, name], flat-prim-app)
+end
+
+var generated-binds = SD.make-mutable-string-dict()
+fun merge-methods(program :: A.Program):
+  doc: ```
+       Tries to merge methods on data definitions where possible
+       Preconditions on program:
+         - well-formed
+         - has been type-checked
+         - contains no s-data
+       Requirements:
+         - all variants have method of the same name
+         - ... with the same arity,
+         - ... and same argument annotations
+       ```
+  cases(A.Program) program block:
+    | s-program(l, _provide, provided-types, imports, body) =>
+      generated-binds := SD.make-mutable-string-dict()
+      { ast: A.s-program(l, _provide, provided-types, imports,
+            if false: body
+            else:
+            body.visit(A.default-map-visitor.{
+                method s-data-expr(self, shadow l, name, name-type, name-ann, params, mixins, variants, shared, _check-loc, _check):
+                  merge-data-methods(l, name, name-type, name-ann, params, mixins, variants, shared, _check-loc, _check)
+                end
+                })
+            end),
+        new-binds: generated-binds }
+    | else => raise("Attempt to desugar non-program: " + torepr(program))
+  end
+end
+fun same-ann(a1 :: A.Ann, a2 :: A.Ann) -> Boolean:
+  a1.visit(A.dummy-loc-visitor) == a2.visit(A.dummy-loc-visitor)
+end
+fun same-sig(s1 :: A.Expr%(is-s-method), s2 :: A.Expr%(is-s-method)):
+  (s1.params.length() == s2.params.length())
+  and for LA.all2-strict(a1 from s1.args, a2 from s2.args):
+    same-ann(a1.ann, a2.ann)
+  end
+  and same-ann(s1.ann, s2.ann)
+end
+fun mk-id-ann(loc, base, ann) block:
+  a = names.make-atom(base)
+  generated-binds.set-now(a.key(), C.value-bind(C.bo-local(loc), C.vb-let, a, ann))
+  { id: a, id-b: A.s-bind(loc, false, a, ann), id-e: A.s-id(loc, a) }
+end
+fun make-renamer():
+  renames = SD.make-mutable-string-dict()
+  fields = SD.make-mutable-string-dict()
+  { renames;
+    fields;
+    A.default-map-visitor.{
+      method s-atom(self, base, serial):
+        n = A.s-atom(base, serial)
+        renames.get-now(n.key()).or-else(n)
+      end,
+      method s-dot(self, l, obj, field) block:
+        cases(A.Expr) obj block:
+          | s-id(_, name) =>
+            # print("Trying to replace " + name.key() + "\n")
+            cases(Option<SD.MutableStringDict<A.Expr>>) fields.get-now(name.key()) block:
+              | some(field-ids) =>
+                cases(Option<A.Expr>) field-ids.get-now(field) block:
+                  | some(id) =>
+                    # print("Replacing " + name.key() + "." + field + " with " + id.id.key() + "\n")
+                    id
+                  | none =>
+                    # print("Couldn't find " + field + " in " + name.key() + ", so recurring\n")
+                    # print(torepr(obj) + "\n")
+                    # print(torepr(obj.visit(self)) + "\n")
+                    # print(torepr(renames.keys-now()) + "\n")
+                    A.s-dot(l, obj.visit(self), field)
+                end
+              | none =>
+                # print("Couldn't find " + name.key() + " at all, so recurring\n")
+                A.s-dot(l, obj.visit(self), field)
+            end
+          | else =>
+            # print("Wasn't a simple s-dot: " + obj.tosource().pretty(10000).first + "." + field + "\n")
+            A.s-dot(l, obj.visit(self), field)
+        end
+      end
+    }
+  }
+end
+fun merge-data-methods(l, name, name-type, name-ann, params, mixins, variants, shared, _check-loc, _check) block:
+  # print("Merging for " + name + "\n")
+  shared-names = SD.make-mutable-string-dict()
+  for each(s from shared) block:
+    shared-names.set-now(s.name, true)
+  end
+  method-sigs = SD.make-mutable-string-dict()
+  method-bodies = SD.make-mutable-string-dict()
+  {renames; fields; renamer} = make-renamer()
+  variants-map = SD.make-mutable-string-dict()
+  needed = variants.length()
+  for each(v from variants) block:
+    variants-map.set-now(v.name, v)
+    for each(w from v.with-members) block:
+      when A.is-s-data-field(w) and A.is-s-method(w.value) block:
+        when not(method-sigs.has-key-now(w.name)) block:
+          method-sigs.set-now(w.name, w.value)
+          method-bodies.set-now(w.name, SD.make-mutable-string-dict())
+        end
+        when same-sig(w.value, method-sigs.get-value-now(w.name)) and not(shared-names.has-key-now(w.name)) block:
+          #print("Candidate for merging type " + name + " : variant " + v.name + " : method " + w.name + "\n")
+          method-bodies.get-value-now(w.name).set-now(v.name, w.value)
+        end
+      end
+    end
+  end
+  shared-methods = SD.make-mutable-string-dict()
+  for SD.each-key-now(m from method-bodies):
+    sig = method-sigs.get-value-now(m)
+    bodies = method-bodies.get-value-now(m)
+    when bodies.count-now() == needed block:
+      # print("Merging type " + name + " : method " + m + "\n")
+      new-params = sig.params.map(lam(n): names.make-atom(n.toname()) end)
+      new-args = sig.args.map(lam(b): mk-id-ann(l, b.id.toname(), b.ann) end)
+      case-bodies = for SD.map-keys-now(vname from bodies) block:
+        vmeth = bodies.get-value-now(vname)
+        for each2(a from vmeth.args, na from new-args):
+          renames.set-now(a.id.key(), na.id)
+        end
+        cases(A.Variant) variants-map.get-value-now(vname) block:
+          | s-variant(lv, constr-loc, _, members, with-members) =>
+            field-exps = SD.make-mutable-string-dict()
+            # print("First arg name for " + vname + " is " + vmeth.args.first.id.key() + "\n")
+            fields.set-now(vmeth.args.first.id.key(), field-exps)
+            arg-fields = members.map(lam(member):
+                cases(A.VariantMember) member block:
+                  | s-variant-member(mloc, mt, b) =>
+                    new-b = mk-id-ann(mloc, b.id.toname(), A.a-blank)
+                    field-exps.set-now(b.id.toname(), new-b.id-e)
+                    cases(A.VariantMemberType) mt:
+                      | s-normal => A.s-cases-bind(mloc, A.s-cases-bind-normal, new-b.id-b)
+                      | s-mutable => A.s-cases-bind(mloc, A.s-cases-bind-ref, new-b.id-b)
+                    end
+                end
+              end)
+            A.s-cases-branch(vmeth.l, constr-loc, vname, arg-fields, vmeth.body.visit(renamer))
+          | s-singleton-variant(lv, _, with-members) =>
+            A.s-singleton-cases-branch(vmeth.l, lv, vname, vmeth.body.visit(renamer))
+        end
+      end
+      else-case = no-method-exn(l, new-args.first.id-e, A.s-str(l, m))
+      ann-name = A.a-name(l, name-ann)
+      shared-method = A.s-data-field(l, m, A.s-method(l, m, new-params, new-args.map(_.id-b), sig.ann, "",
+        A.s-cases-else(l, ann-name, new-args.first.id-e, case-bodies, else-case, true),
+        none, none, true))
+      shared-methods.set-now(m, shared-method)
+    end
+  end
+  new-variants = for map(v from variants):
+    cases(A.Variant) v block:
+      | s-variant(lv, constr-loc, vname, members, with-members) =>
+        A.s-variant(lv, constr-loc, vname, members,
+          with-members.filter(lam(m): not(shared-methods.has-key-now(m.name)) end))
+      | s-singleton-variant(lv, vname, with-members) =>
+        A.s-singleton-variant(lv, vname, with-members.filter(lam(m): not(shared-methods.has-key-now(m.name)) end))
+    end
+  end
+  A.s-data-expr(l, name, name-type, name-ann, params, mixins,
+    new-variants,
+    shared-methods.map-keys-now(lam(mname): shared-methods.get-value-now(mname) end) + shared,
+    _check-loc, _check)
+end
+
 fun desugar-post-tc(program :: A.Program, compile-env :: C.CompileEnvironment):
   doc: ```
         Desugar non-scope and non-check based constructs.

src/arr/compiler/desugar.arr

@@ -387,7 +387,7 @@ fun desugar-expr(expr :: A.Expr):
       A.s-let-expr(l, new-binds, desugar-expr(body), blocky)
     | s-letrec(l, binds, body, blocky) =>
       A.s-letrec(l, desugar-letrec-binds(binds), desugar-expr(body), blocky)
-    | s-data-expr(l, name, namet, params, mixins, variants, shared, _check-loc, _check) =>
+    | s-data-expr(l, name, name-type, name-ann, params, mixins, variants, shared, _check-loc, _check) =>
       fun extend-variant(v):
         cases(A.Variant) v:
           | s-variant(l2, constr-loc, vname, members, with-members) =>
@@ -404,7 +404,7 @@ fun desugar-expr(expr :: A.Expr):
               with-members.map(desugar-member))
         end
       end
-      A.s-data-expr(l, name, namet, params, mixins.map(desugar-expr), variants.map(extend-variant),
+      A.s-data-expr(l, name, name-type, name-ann, params, mixins.map(desugar-expr), variants.map(extend-variant),
         shared.map(desugar-member), _check-loc, desugar-opt(desugar-expr, _check))
     | s-when(l, test, body, blocky) =>
       ds-test = desugar-expr(test)

src/arr/compiler/resolve-scope.arr

@@ -102,9 +102,10 @@ fun desugar-toplevel-types(stmts :: List<A.Expr>) -> List<A.Expr> block:
       | s-newtype(l, name, namet) =>
         rev-type-binds := link(A.s-newtype-bind(l, name, namet), rev-type-binds)
       | s-data(l, name, params, mixins, variants, shared, _check-loc, _check) =>
-        namet = names.make-atom(name)
-        rev-type-binds := link(A.s-newtype-bind(l, A.s-name(l, name), namet), rev-type-binds)
-        rev-stmts := link(A.s-data-expr(l, name, namet, params, mixins, variants, shared, _check-loc, _check), rev-stmts)
+        name-type = names.make-atom(name)
+        name-ann = A.s-name(l, name) # placeholder until name resolution happens
+        rev-type-binds := link(A.s-newtype-bind(l, A.s-name(l, name), name-type), rev-type-binds)
+        rev-stmts := link(A.s-data-expr(l, name, name-type, name-ann, params, mixins, variants, shared, _check-loc, _check), rev-stmts)
       | else =>
         rev-stmts := link(s, rev-stmts)
     end
@@ -420,7 +421,7 @@ fun desugar-scope-block(stmts :: List<A.Expr>, binding-group :: BindingGroup) ->
               # it'll get turned into an s-lam in weave-contracts
               f
               ), rest-stmts)
-        | s-data-expr(l, name, namet, params, mixins, variants, shared, _check-loc, _check) =>
+        | s-data-expr(l, name, name-type, name-ann, params, mixins, variants, shared, _check-loc, _check) =>
           fun b(loc, id :: String): A.s-bind(loc, false, A.s-name(loc, id), A.a-blank) end
           fun bn(loc, n :: A.Name): A.s-bind(loc, false, n, A.a-blank) end
           fun variant-binds(data-blob-id, v):
@@ -433,7 +434,7 @@ fun desugar-scope-block(stmts :: List<A.Expr>, binding-group :: BindingGroup) ->
             ]
           end
           blob-id = names.make-atom(name)
-          data-expr = A.s-data-expr(l, name, namet, params, mixins, variants, shared, _check-loc, _check)
+          data-expr = A.s-data-expr(l, name, name-type, name-ann, params, mixins, variants, shared, _check-loc, _check)
           bind-data = A.s-letrec-bind(l, bn(l, blob-id), data-expr)
           bind-data-pred = A.s-letrec-bind(l, b(l, A.make-checker-name(name)), A.s-dot(l, A.s-id-letrec(l, blob-id, true), name))
           all-binds = for fold(acc from [list: bind-data-pred, bind-data], v from variants):
@@ -1059,7 +1060,7 @@ fun resolve-names(p :: A.Program, initial-env :: C.CompileEnvironment):
       data-defs = for lists.filter-map(ddk from datatypes.keys-list-now()):
         dd = datatypes.get-value-now(ddk) 
         if provide-types-dict.has-key(dd.name):
-          some(A.p-data(dd.l, dd.namet, none))
+          some(A.p-data(dd.l, dd.name-type, none))
         else:
           none
         end
@@ -1187,18 +1188,19 @@ fun resolve-names(p :: A.Program, initial-env :: C.CompileEnvironment):
       A.s-cases-branch(l, pat-loc, name, new-args, new-body)
     end,
     # s-singleton-cases-branch introduces no new bindings
-    method s-data-expr(self, l, name, namet, params, mixins, variants, shared-members, _check-loc, _check) block:
+    method s-data-expr(self, l, name, name-type, _, params, mixins, variants, shared-members, _check-loc, _check) block:
       {env; atoms} = for fold(acc from { self.type-env; empty }, param from params):
         {env; atoms} = acc
         atom-env = make-atom-for(param, false, env, type-bindings,
           C.type-bind(C.bo-local(l), C.tb-type-var, _, none))
         { atom-env.env; link(atom-env.atom, atoms) }
       end
       with-params = self.{type-env: env}
-      result = A.s-data-expr(l, name, namet, atoms.reverse(),
+      name-ann = self.type-env.get-value(name).atom
+      result = A.s-data-expr(l, name, name-type, name-ann, atoms.reverse(),
         mixins.map(_.visit(with-params)), variants.map(_.visit(with-params)),
         shared-members.map(_.visit(with-params)), _check-loc, with-params.option(_check))
-      datatypes.set-now(namet.key(), result)
+      datatypes.set-now(name-type.key(), result)
       result
     end,
     method s-lam(self, l, name, params, args, ann, doc, body, _check-loc, _check, blocky) block: