fixup: make more things work for the trie refactor

docs/src/concepts/theory_composition.md

@@ -0,0 +1,51 @@
+# Theory Composition
+
+As theories get larger, it becomes more and more important to not build the
+entire theory from scratch. Not only is this tedious, it is also error-prone.
+From the beginning, Catlab and GATlab have supported single inheritance, which
+helps to some extent. In this document, we lay out other approaches to composing
+theories.
+
+## Multiple Inheritance
+
+In a GATlab `@theory`, one can use `using` to take the *union* of one theory
+with another theory.
+
+The way this works is the following. Every time a new theory is created, the new
+definitions for that theory form a new scope, with a unique UUID. Union of
+theories operates on a scope tag level, taking the union of the sets of UUIDs
+and then producing a theory with all the bindings from the scopes tagged by
+those UUIDs.
+
+If we never had to parse user-supplied expressions, then the names of the
+operations in the theories wouldn't matter, because identifiers come with scope
+tags. However, as a practical matter, we disallow unioning two theories with the
+same name declaration.
+
+That being said, it is fine to union two theories which *overload* the same
+declaration. That is, if two theories have the declaration of a name in common,
+then they can overload that name as long as they don't give conflicting
+overloads, in the same way that overloading methods in Julia works.
+
+This is akin to the way multiple inheritance works in frameworks such as
+
+- Haskell typeclasses
+- [Object-oriented systems with multiple inheritance, like Scala](https://docs.scala-lang.org/scala3/book/domain-modeling-tools.html#traits)
+- [Module inclusion in OCaml](https://cs3110.github.io/textbook/chapters/modules/includes.html)
+
+## Nesting
+
+However, there are other ways of composing things like GATlab theories. In
+dependently typed languages used for theorem proving, algebraic structures are
+often represented by dependent records. For instance, in the agda unimath
+library, the [definition of a group](https://github.com/UniMath/agda-unimath/blob/master/src/group-theory/groups.lagda.md) is
+
+```agda
+Semigroup :
+  (l : Level) → UU (lsuc l)
+Semigroup l = Σ (Set l) has-associative-mul-Set
+
+Group :
+  (l : Level) → UU (lsuc l)
+Group l = Σ (Semigroup l) is-group
+```

docs/src/examples/springs.md

@@ -0,0 +1,5 @@
+# Composition of resource sharers
+
+```julia
+
+```

src/nonstdlib/dynamics/ResourceSharers.jl

@@ -1,10 +1,11 @@
 module ResourceSharers
-export Rhizome, @rhizome, ResourceSharer, Variable
+export Rhizome, @rhizome, ResourceSharer, @resource_sharer, Variable
 
 using ...Syntax
 using ...Syntax.Tries
 using ...Syntax.Tries: flatten, node
 using ...Syntax.GATs: tcompose
+using ...Util
 using MLStyle
 using OrderedCollections
 
@@ -77,11 +78,11 @@ The only problem is that I'm not sure then what to call the single field in
 `Rhizome`. Also, this would require a slight refactor of the rest of the code.
 """
 struct Rhizome
+  theory::GAT
+  mcompose::AlgClosure
+  mzero::AlgClosure
   boxes::Trie{Trie{PortVariable}}
   junctions::Trie{Variable}
-  function Rhizome(boxes::Trie{Trie{PortVariable}}, junctions::Trie{Variable})
-    new(boxes, junctions)
-  end
 end
 
 """
@@ -102,15 +103,15 @@ databases, recursion, and plug-and-play circuits](https://arxiv.org/abs/1305.029
 function ocompose(r::Rhizome, rs::Trie{Rhizome})
   # paired :: Trie{Tuple{Trie{PortVariable}, Rhizome}}
   paired = zip(r.boxes, rs)
-  boxes = flatten(mapwithkey(paired) do k, (interface, r′)
+  boxes = flatten(mapwithkey(Trie{Trie{PortVariable}}, paired) do k, (interface, r′)
     # k :: TrieVar
     # interface :: Trie{PortVariable}
     # r′ :: Rhizome
     # We want to create the new collection of boxes
 
-    map(r′.boxes) do b
+    map(Trie{PortVariable}, r′.boxes) do b
       # b :: Trie{PortVariable}
-      map(b) do p
+      map(PortVariable, b) do p
         # p :: PortVariable
         # p.junction :: namespace(b.junctions)
         jvar = p.junction
@@ -128,7 +129,7 @@ function ocompose(r::Rhizome, rs::Trie{Rhizome})
   end)
   # Add all unexposed junctions
   newjunctions = flatten(
-    map(rs) do r′
+    map(Trie{Variable}, rs) do r′
       internal_junctions = filter(j -> !j.exposed, r′.junctions)
       if isnothing(internal_junctions)
         Tries.node(OrderedDict{Symbol, Trie{Variable}}())
@@ -140,7 +141,7 @@ function ocompose(r::Rhizome, rs::Trie{Rhizome})
 
   junctions = merge(r.junctions, newjunctions)
 
-  Rhizome(boxes, junctions)
+  Rhizome(r.theory, r.mcompose, r.mzero, boxes, junctions)
 end
 
 # TODO: this should just use a `toexpr` method
@@ -210,9 +211,9 @@ will be assumed to be of that type.
 3. The namespace for the external ports of the rhizome is `[a, b]`.
 As noted in 1., each of these ports is typed with `default`.
 """
-macro rhizome(theory, head, body)
+macro rhizome(theorymod, head, body)
   # macroexpand `theory` to get the actual theory
-  theory = macroexpand(__module__, :($theory.Meta.@theory))
+  theory = macroexpand(__module__, :($theorymod.Meta.@theory))
   # parse the name and junctions out of `head`
   junctions = OrderedDict{TrieVar, Variable}()
   (name, args) = @match head begin
@@ -254,8 +255,14 @@ macro rhizome(theory, head, body)
     boxes[box] = Trie(interface)
   end
   :(
-    $(esc(name)) = $(Rhizome(Trie(boxes), junctions))
-  )
+    $name = $Rhizome(
+      $theory,
+      $theorymod.Meta.Constructors.:(+),
+      $theorymod.Meta.Constructors.zero,
+      $(Trie(boxes)),
+      $(junctions),
+    )
+  ) |> esc
 end
 
 """
@@ -275,6 +282,82 @@ struct ResourceSharer
   # observe::AlgClosure
 end
 
+"""
+A DSL for writing down resource sharers
+
+Example:
+```
+@resource_sharer ThRing Spring begin
+  variables = x, v
+  params = k
+  update = (state, params) -> (x = state.v, v = -params.k * state.x)
+end
+```
+"""
+macro resource_sharer(theory, name, body)
+  args = Expr[]
+
+  for line in body.args
+    @match line begin
+      _ :: LineNumberNode => nothing
+      Expr(:(=), arg, val) =>
+        push!(args, Expr(:kw, arg, Expr(:quote, val)))
+    end
+  end
+
+  esc(:($name = $ResourceSharer($theory.Meta.theory; $(args...))))
+end
+
+function parse_namespace(theory::GAT, expr::Expr0)
+  vs = OrderedDict{TrieVar, AlgType}()
+  argexprs = @match expr begin
+    Expr(:tuple, args...) => args
+    _ => [expr]
+  end
+  for arg in argexprs
+    (vname, typeexpr) = @match arg begin
+      vname::Symbol => (vname, :default)
+      Expr(:(.), _, _) => (arg, :default)
+      Expr(:(::), vname, type) => (vname, type)
+    end
+    v = parse_var(vname)
+    type = fromexpr(theory, typeexpr, AlgType)
+    vs[v] = type
+  end
+  Trie(vs)
+end
+
+function ResourceSharer(theory::GAT; variables::Expr0, params::Expr0, update::Expr0)
+  variable_trie = parse_namespace(theory, variables)
+  variables = map(Variable, variable_trie) do type
+    Variable(true, type)
+  end
+  param_type = tcompose(parse_namespace(theory, params))
+  state_type = tcompose(variable_trie)
+
+  update = begin
+    (argexpr, bodyexpr) = @match update begin
+      Expr(:(->), args, body) => (args, last(body.args))
+    end
+    statename, paramname = @match argexpr begin
+      Expr(:tuple, s, p) => (s, p)
+    end
+    args = TypeScope(statename => state_type, paramname => param_type)
+    body = fromexpr(GATContext(theory, args), bodyexpr, AlgTerm)
+    AlgMethod(args, body, "", LID.([1,2]), state_type)
+  end
+
+  ResourceSharer(variables, param_type, update)
+end
+
+function Base.show(io::IO, r::ResourceSharer; theory)
+  println(io, "ResourceSharer:")
+  print(io, "variables = ", map(v -> string(toexpr(theory, v.type)), String, r.variables))
+  println(io, "params = ", toexpr(theory, r.params))
+  println(io, "update = ")
+  show(io, r.update; theory)
+end
+
 """
 If C is a cartesian category, and t1 and t2 are families of objects in C with a function
 of finite sets dom(t2) -> dom(t1), then we can use "copy and delete" to form a morphism in C
@@ -294,7 +377,7 @@ function pullback(t1::Trie{AlgType}, t2::Trie{Tuple{AlgType, TrieVar}}; argname=
   ctx = TypeScope(argname => ty1)
   x = AlgTerm(ident(ctx; name=argname))
   body = tcompose(
-    map(t2) do (_, k)
+    map(AlgTerm, t2) do (_, k)
       x[k]
     end
   )
@@ -318,8 +401,8 @@ Produces an AlgMethod going from tcompose(first.(t2)) to tcompose(t1)
 function pushforward(
   t1::Trie{AlgType},
   t2::Trie{Tuple{AlgType, TrieVar}},
-  mcompose::Tuple{Ident, Ident},
-  mzero::Tuple{Ident, Ident};
+  mcompose::AlgClosure,
+  mzero::AlgClosure;
   argname=:x
 )
   preimages = Dict{TrieVar, Vector{TrieVar}}()
@@ -335,18 +418,18 @@ function pushforward(
   ctx = TypeScope(argname => ty2)
   x = AlgTerm(ident(ctx; name=argname))
   body = tcompose(
-    mapwithkey(t1) do k, _
+    mapwithkey(AlgTerm, t1) do k, _
       foldl(
-        (term, v) -> AlgTerm(mcompose..., [term, x[v]]),
+        (term, v) -> first(values(mcompose.methods))(term, x[v]),
         preimages[k];
-        init=AlgTerm(mzero..., AlgTerm[])
+        init=mzero()
       )
     end
   )
   AlgMethod(ctx, body, "", [LID(1)], ty1)
 end
 
-function oapply(r::Rhizome, sharers::Trie{ResourceSharer}, mcompose, mzero)
+function oapply(r::Rhizome, sharers::Trie{ResourceSharer})
   new_variables = filter(v -> !v.exposed, flatten(
     map(sharers) do sharer
       sharer.variables
@@ -365,8 +448,8 @@ function oapply(r::Rhizome, sharers::Trie{ResourceSharer}, mcompose, mzero)
   # to their type, and to the variable in the reduced system that they
   # map to
   full_state = flatten(
-    mapwithkey(zip(r.boxes, sharers)) do b, (interface, sharer)
-      mapwithkey(sharer.variables) do k, v
+    mapwithkey(Trie{Tuple{AlgType, TrieVar}}, zip(r.boxes, sharers)) do b, (interface, sharer)
+      mapwithkey(Tuple{AlgType, TrieVar}, sharer.variables) do k, v
         if v.exposed
           (v.type, interface[k].junction)
         else
@@ -385,7 +468,7 @@ function oapply(r::Rhizome, sharers::Trie{ResourceSharer}, mcompose, mzero)
   copy = pullback(state, full_state; argname=:state)
 
   # add :: full_state -> state
-  add = pushforward(state, full_state, mcompose, mzero)
+  add = pushforward(state, full_state, r.mcompose, r.mzero)
 
   update_ctx = TypeScope(:state => state_type, :params => params)
   statevar, paramsvar = AlgTerm.(idents(update_ctx; name=[:state, :params]))

src/syntax/Tries.jl

@@ -121,6 +121,9 @@ function content(p::Trie)
   end
 end
 
+"""
+Construct a new Trie node.
+"""
 function node(d::OrderedDict{Symbol, <:AbstractTrie{A}}) where {A}
   nonempties = OrderedDict{Symbol, NonEmptyTrie{A}}()
   for (k, t) in pairs(d)
@@ -134,7 +137,7 @@ function node(d::OrderedDict{Symbol, <:AbstractTrie{A}}) where {A}
   if length(nonempties) > 0
     Trie{A}(NonEmptyTrie{A}(Node_{NonEmptyTrie{A}}(nonempties)))
   else
-    Trie{A}(nothing)
+    Trie{A}()
   end
 end
 
@@ -267,7 +270,12 @@ function Base.map(f, return_type::Type, t::AbstractTrie)
   end
 end
 
-function flatten(t::Trie{Trie{A}})::Trie{A} where {A}
+function Base.map(f, t::AbstractTrie{A}) where {A}
+  B = Core.Compiler.return_type(f, Tuple{A})
+  map(f, B, t)
+end
+
+function flatten(t::AbstractTrie{Trie{A}}) where {A}
   @match t begin
     Leaf(v) => v
     Node(bs) => node(n => flatten(v) for (n, v) in bs)
@@ -315,16 +323,16 @@ function mapwithkey(f, return_type::Type, t::AbstractTrie; prefix=PACKAGE_ROOT)
   @match t begin
     Leaf(v) => leaf(f(prefix, v))
     Node(bs) =>
-      node([k => mapwithkey(f, v; prefix=getproperty(prefix, k)) for (k, v) in bs]...)
+      node([k => mapwithkey(f, return_type, v; prefix=getproperty(prefix, k)) for (k, v) in bs]...)
     Empty() => Trie{return_type}()
   end
 end
 
-function traversewithkey(f, t::Trie; prefix=PACKAGE_ROOT)
+function traversewithkey(f, t::AbstractTrie; prefix=PACKAGE_ROOT)
   @match t begin
     Leaf(v) => (f(prefix, v); nothing)
     Node(bs) => begin
-      for (k, v) in branches(t)
+      for (k, v) in bs
         traversewithkey(f, v; prefix=getproperty(prefix, k))
       end
     end
@@ -349,8 +357,8 @@ function Base.merge(t1::AbstractTrie{A}, t2::AbstractTrie{A}) where {A}
   @match (t1, t2) begin
     (Leaf(_), _) || (_, Leaf(_)) =>
       error("cannot merge tries with overlapping keys")
-    (Empty, _) => t2
-    (_, Empty) => t1
+    (Empty(), _) => t2
+    (_, Empty()) => t1
     (Node(b1), Node(b2)) => begin
       b = OrderedDict{Symbol, NonEmptyTrie{A}}()
       for (n, t) in b1
@@ -378,19 +386,19 @@ Fails if the keys in the dict are not prefix-free.
 function Trie(d::OrderedDict{TrieVar, A}) where {A}
   branches = OrderedDict{Symbol, OrderedDict{TrieVar, A}}()
   for (v, x) in d
-    if isroot(v)
-      if length(d) == 1
-        return leaf(x)
-      else
-        error("attempted trie conversion failed because keys were not prefix-free")
-      end
-    else
-      (n, v′) = content(v)
-      if haskey(branches, n)
-        branches[n][v′] = x
-      else
-        branches[n] = OrderedDict(v′ => x)
-      end
+    @match v begin
+      Root() =>
+        if length(d) == 1
+          return leaf(x)
+        else
+          error("attempted trie conversion failed because keys were not prefix-free")
+        end
+      Nested((n, v′)) =>
+        if haskey(branches, n)
+          branches[n][v′] = x
+        else
+          branches[n] = OrderedDict(v′ => x)
+        end
     end
   end
   node(n => Trie(d′) for (n, d′) in branches)