ENH: Allow GAT expressions in subpart and incident.

AlgebraicJulia · Oct 19, 2020 · 50d9b78 · 50d9b78
1 parent d125e5d
commit 50d9b78
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/src/categorical_algebra/CSetDataStructures.jl b/src/categorical_algebra/CSetDataStructures.jl
@@ -12,6 +12,7 @@ using Compat: isnothing, only
 using PrettyTables: pretty_table
 using StructArrays
 
+using ...Syntax: GATExpr, args
 using ...Theories: Schema, FreeSchema, dom, codom, codom_num,
   CatDesc, CatDescType, AttrDesc, AttrDescType, SchemaType
 using ...Present
@@ -340,6 +341,11 @@ function subpart(acs::ACSet, part, names::AbstractVector{Symbol})
     subpart(acs, part, name)
   end
 end
+subpart(acs::ACSet, part, expr::GATExpr) = subpart(acs, part, subpart_name(expr))
+
+subpart_name(expr::GATExpr{:generator}) = first(expr)::Symbol
+subpart_name(expr::GATExpr{:id}) = Symbol[]
+subpart_name(expr::GATExpr{:compose}) = mapreduce(subpart_name, vcat, args(expr))
 
 @generated function _subpart(acs::ACSet{CD,AD,Ts}, ::Val{name}) where
     {CD,AD,Ts,name}
@@ -373,6 +379,8 @@ function incident(acs::ACSet, part, names::AbstractVector{Symbol};
     reduce(vcat, incident(acs, part, name), init=Int[])
   end
 end
+incident(acs::ACSet, part, expr::GATExpr; kw...) =
+  incident(acs, part, subpart_name(expr); kw...)
 
 @generated function _incident(acs::ACSet{CD,AD,Ts,Idxed}, part, ::Val{name};
                               copy::Bool=false) where {CD,AD,Ts,Idxed,name}

diff --git a/test/categorical_algebra/CSetDataStructures.jl b/test/categorical_algebra/CSetDataStructures.jl
@@ -1,7 +1,8 @@
 module TestCSetDataStructures
 using Test
 
-using Catlab: @present
+using Catlab: @present, generator
+using Catlab.Theories: compose, id
 using Catlab.CSetDataStructures
 
 # Discrete dynamical systems
@@ -151,6 +152,12 @@ set_subpart!(d, [4,5], :parent, 5)
 @test incident(d, 5, [:parent, :parent]) == [1,2,3,4,5]
 @test incident(d, 10, [:parent, :height]) == [1,2,3]
 
+X, parent, height = generator.(Ref(TheoryDendrogram), [:X, :parent, :height])
+@test subpart(d, 3, parent) == 4
+@test subpart(d, 3, compose(parent, height)) == 10
+@test subpart(d, 3, id(X)) == 3
+@test incident(d, 10, compose(parent, height)) == [1,2,3]
+
 # Copying parts.
 d2 = Dendrogram{Int}()
 copy_parts!(d2, d, X=[4,5])