Merge pull request #36 from AlgebraicJulia/index_collect

Collect by default, opt-in for view of column data
AlgebraicJulia · Jun 21, 2023 · c90024d · c90024d
2 parents 869d2a2 + 682b2f7
commit c90024d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 6 deletions.
diff --git a/src/ACSetInterface.jl b/src/ACSetInterface.jl
@@ -102,9 +102,6 @@ function subpart(acs, part, names::AbstractVector{Symbol})
   end
 end
 
-subpart(acs, names::AbstractVector{Symbol}) =
-  subpart(acs, Int[subpart(acs, names[1])...], names[2:end])
-
 Base.getindex(acs::ACSet, part, name) = subpart(acs, part, name)
 Base.getindex(acs::ACSet, name) = subpart(acs, name)
 

diff --git a/src/DenseACSets.jl b/src/DenseACSets.jl
@@ -376,13 +376,30 @@ end
   end
 end
 
-@inline ACSetInterface.subpart(acs::SimpleACSet, f::Symbol) =
+Base.view(acs::SimpleACSet, part, f) = view_with_default(acs.subparts[f], part, default_value(acs, f))
+Base.view(acs::SimpleACSet, part::AbstractVector{Bool}, f) = view(acs, f)[part]
+Base.view(acs::SimpleACSet, ::Colon, f) = view(acs, dom_parts(acs,f), f)
+
+@inline Base.view(acs::SimpleACSet, f::Symbol) =
   view_with_default(acs.subparts[f], dom_parts(acs, f), default_value(acs, f))
 
+@inline ACSetInterface.subpart(acs::SimpleACSet, f::Symbol) = subpart(acs, dom_parts(acs, f), f)
+@inline ACSetInterface.subpart(acs::SimpleACSet, f::Vector{Symbol}) = subpart(acs, dom_parts(acs, first(f)), f)
+@inline ACSetInterface.subpart(acs::ACSet, part::Union{Colon,AbstractVector}, name::Symbol) =
+  collect_column(view(acs, part, name))
+
+function collect_column(x::AbstractVector)
+  if isempty(x)
+    Base.typesplit(eltype(x), AttrVar)[]
+  else
+    map(identity, x)
+  end
+end
+
 @inline ACSetInterface.subpart(acs::SimpleACSet, part::Int, f::Symbol) =
   get(acs.subparts[f], part, default_value(acs, f))
 
-@inline ACSetInterface.has_subpart(acs::StructACSet{S}, f::Symbol) where {S} =
+@inline ACSetInterface.has_subpart(::StructACSet{S}, f::Symbol) where {S} =
   _has_subpart(Val{S}, Val{f})
 
 ACSetInterface.has_subpart(acs::DynamicACSet, f::Symbol) =

diff --git a/test/ACSets.jl b/test/ACSets.jl
@@ -5,6 +5,7 @@ using StaticArrays: StaticVector
 using Tables
 
 using ACSets
+using ACSets.Columns: ColumnView
 
 # Discrete dynamical systems
 ############################
@@ -42,6 +43,7 @@ for dds_maker in dds_makers
   @test add_part!(dds, :X, Φ=1) == 3
   @test subpart(dds, :Φ) == [1,1,1]
   @test subpart(dds, [2,3], :Φ) == [1,1]
+  @test subpart(dds, Bool[0,1,1], :Φ) == [1,1]
   @test incident(dds, 1, :Φ) == [1,2,3]
 
   @test has_part(dds, :X)
@@ -56,9 +58,17 @@ for dds_maker in dds_makers
   @test_throws Exception incident(dds, 1, :nonsuppart)
   @test_throws Exception set_subpart!(dds, 1, :nonsubpart, 1)
 
-  # Deletion.
+  # Types of subset-ed columns.
   dds = dds_maker()
   add_parts!(dds, :X, 3, Φ=[2,3,3])
+  dds[1:2,:Φ] isa Vector{Int}
+  dds[:Φ] isa Vector{Int}
+  dds[[:Φ,:Φ]] isa Vector{Int}
+  dds[1:2, [:Φ,:Φ]] isa Vector{Int}
+  view(dds,1:2,:Φ) isa ColumnView
+  view(dds,:Φ) isa ColumnView
+
+  # Deletion.
   rem_part!(dds, :X, 2)
   @test nparts(dds, :X) == 2
   @test subpart(dds, :Φ) == [0,2]