More reorg

Author: mtfishman

src/lazynameddimsarrays.jl

@@ -37,8 +37,8 @@ end
 
 # Custom version of `AbstractTrees.printnode` to
 # avoid type piracy when overloading on `AbstractNamedDimsArray`.
-printnode(io::IO, x) = AbstractTrees.printnode(io, x)
-function printnode(io::IO, a::AbstractNamedDimsArray)
+printnode_nameddims(io::IO, x) = AbstractTrees.printnode(io, x)
+function printnode_nameddims(io::IO, a::AbstractNamedDimsArray)
     show(io, collect(dimnames(a)))
     return nothing
 end
@@ -149,6 +149,37 @@ function map_arguments_lazy(f, a)
         return error("Variant not supported.")
     end
 end
+function substitute_lazy(a, substitutions::AbstractDict)
+    haskey(substitutions, a) && return substitutions[a]
+    !iscall(a) && return a
+    return map_arguments(arg -> substitute(arg, substitutions), a)
+end
+function substitute_lazy(a, substitutions)
+    return substitute(a, Dict(substitutions))
+end
+function printnode_lazy(io, a)
+    # Use `printnode_nameddims` to avoid type piracy,
+    # since it overloads on `AbstractNamedDimsArray`.
+    return printnode_nameddims(io, unwrap(a))
+end
+function show_lazy(io::IO, a)
+    if !iscall(a)
+        return show(io, unwrap(a))
+    else
+        return AbstractTrees.printnode(io, a)
+    end
+end
+function show_lazy(io::IO, mime::MIME"text/plain", a)
+    summary(io, a)
+    println(io, ":")
+    if !iscall(a)
+        show(io, mime, unwrap(a))
+        return nothing
+    else
+        show(io, a)
+        return nothing
+    end
+end
 
 # Lazy broadcasting.
 struct LazyNamedDimsArrayStyle <: AbstractNamedDimsArrayStyle{Any} end
@@ -185,7 +216,7 @@ TermInterface.head(m::Applied) = operation(m)
 TermInterface.iscall(m::Applied) = true
 TermInterface.isexpr(m::Applied) = iscall(m)
 function Base.show(io::IO, m::Applied)
-    args = map(arg -> sprint(printnode, arg), arguments(m))
+    args = map(arg -> sprint(AbstractTrees.printnode, arg), arguments(m))
     print(io, "(", join(args, " $(operation(m)) "), ")")
     return nothing
 end
@@ -301,37 +332,20 @@ end
 function map_arguments(f, a::LazyNamedDimsArray)
     return map_arguments_lazy(f, a)
 end
-
-function substitute(a::LazyNamedDimsArray, substitutions::AbstractDict)
-    haskey(substitutions, a) && return substitutions[a]
-    !iscall(a) && return a
-    return map_arguments(arg -> substitute(arg, substitutions), a)
-end
 function substitute(a::LazyNamedDimsArray, substitutions)
-    return substitute(a, Dict(substitutions))
-end
-
-function printnode(io::IO, a::LazyNamedDimsArray)
-    return printnode(io, unwrap(a))
+    return substitute_lazy(a, substitutions)
 end
 function AbstractTrees.printnode(io::IO, a::LazyNamedDimsArray)
-    return printnode(io, a)
+    return printnode_lazy(io, a)
+end
+function printnode_nameddims(io::IO, a::LazyNamedDimsArray)
+    return printnode_lazy(io, a)
 end
 function Base.show(io::IO, a::LazyNamedDimsArray)
-    if !iscall(a)
-        return show(io, unwrap(a))
-    else
-        return printnode(io, a)
-    end
+    return show_lazy(io, a)
 end
 function Base.show(io::IO, mime::MIME"text/plain", a::LazyNamedDimsArray)
-    if !iscall(a)
-        @invoke show(io, mime, a::AbstractNamedDimsArray)
-        return nothing
-    else
-        show(io, a)
-        return nothing
-    end
+    return show_lazy(io, mime, a)
 end
 
 function Base.:*(a::LazyNamedDimsArray)
@@ -427,13 +441,16 @@ const SymbolicNamedDimsArray{T, N, Parent <: SymbolicArray{T, N}, DimNames} =
 function symnameddims(name)
     return lazy(NamedDimsArray(SymbolicArray(name), ()))
 end
-function printnode(io::IO, a::SymbolicNamedDimsArray)
+function AbstractTrees.printnode(io::IO, a::SymbolicNamedDimsArray)
     print(io, symname(dename(a)))
     if ndims(a) > 0
         print(io, "[", join(dimnames(a), ","), "]")
     end
     return nothing
 end
+function printnode_nameddims(io::IO, a::SymbolicNamedDimsArray)
+    return AbstractTrees.printnode(io, a)
+end
 function Base.:(==)(a::SymbolicNamedDimsArray, b::SymbolicNamedDimsArray)
     return issetequal(inds(a), inds(b)) && dename(a) == dename(b)
 end

test/test_lazynameddimsarrays.jl

@@ -2,7 +2,7 @@ using AbstractTrees: AbstractTrees, print_tree, printnode
 using Base.Broadcast: materialize
 using ITensorNetworksNext.LazyNamedDimsArrays:
     LazyNamedDimsArray, Mul, SymbolicArray, ismul, lazy, substitute, symnameddims
-using NamedDimsArrays: NamedDimsArray, @names, dename, dimnames, inds, nameddims
+using NamedDimsArrays: NamedDimsArray, @names, dename, dimnames, inds, nameddims, namedoneto
 using TermInterface:
     arguments,
     arity,
@@ -19,9 +19,10 @@ using WrappedUnions: unwrap
 
 @testset "LazyNamedDimsArrays" begin
     @testset "Basics" begin
-        a1 = nameddims(randn(2, 2), (:i, :j))
-        a2 = nameddims(randn(2, 2), (:j, :k))
-        a3 = nameddims(randn(2, 2), (:k, :l))
+        i, j, k, l = namedoneto.(2, (:i, :j, :k, :l))
+        a1 = randn(i, j)
+        a2 = randn(j, k)
+        a3 = randn(k, l)
         l1, l2, l3 = lazy.((a1, a2, a3))
         for li in (l1, l2, l3)
             @test li isa LazyNamedDimsArray
@@ -81,7 +82,8 @@ using WrappedUnions: unwrap
         @test AbstractTrees.children(l) == [l1 * l2, l3]
         @test AbstractTrees.nodevalue(l) ≡ *
         @test sprint(show, l) == "(([:i, :j] * [:j, :k]) * [:k, :l])"
-        @test sprint(show, MIME"text/plain"(), l) == "(([:i, :j] * [:j, :k]) * [:k, :l])"
+        @test sprint(show, MIME"text/plain"(), l) ==
+            "named(Base.OneTo(2), :i)×named(Base.OneTo(2), :l) LazyNamedDimsArray{Float64, …}:\n(([:i, :j] * [:j, :k]) * [:k, :l])"
         @test sprint(printnode, l) == "(([:i, :j] * [:j, :k]) * [:k, :l])"
         @test sprint(print_tree, l) ==
             "(([:i, :j] * [:j, :k]) * [:k, :l])\n├─ ([:i, :j] * [:j, :k])\n│  ├─ [:i, :j]\n│  └─ [:j, :k]\n└─ [:k, :l]\n"
@@ -102,7 +104,8 @@ using WrappedUnions: unwrap
         @test arguments(ex) == [a1 * a2, a3]
         @test operation(ex) ≡ *
         @test sprint(show, ex) == "((a1 * a2) * a3)"
-        @test sprint(show, MIME"text/plain"(), ex) == "((a1 * a2) * a3)"
+        @test sprint(show, MIME"text/plain"(), ex) ==
+            "0-dimensional LazyNamedDimsArray{Any, …}:\n((a1 * a2) * a3)"
     end
 
     @testset "substitute" begin