From 7aad2971fc5f193390b66c36f503ed1dcecc087d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Olivier=20Gauth=C3=A9?= <ogauthe@flatironinstitute.org>
Date: Wed, 15 Jan 2025 16:17:58 -0500
Subject: [PATCH] improve names and loops

---
 src/fusion_trees/fusiontree.jl   | 49 +++++++++++++++++---------------
 src/fusiontensor/fusiontensor.jl | 20 ++++++-------
 2 files changed, 35 insertions(+), 34 deletions(-)

diff --git a/src/fusion_trees/fusiontree.jl b/src/fusion_trees/fusiontree.jl
index 2f48757..bc7ad38 100644
--- a/src/fusion_trees/fusiontree.jl
+++ b/src/fusion_trees/fusiontree.jl
@@ -81,15 +81,6 @@ Base.length(::SectorFusionTree{<:Any,N}) where {N} = N
 # GradedUnitRanges interface
 GradedUnitRanges.sector_type(::Type{<:SectorFusionTree{S}}) where {S} = S
 
-function build_trees(legs::Vararg{AbstractGradedUnitRange})
-  tree_arrows = isdual.(legs)
-  sectors = blocklabels.(legs)
-  return mapreduce(vcat, CartesianIndices(blocklength.(legs))) do it
-    block_sectors = getindex.(sectors, Tuple(it))  # why not type stable?
-    return build_trees(block_sectors, tree_arrows)
-  end
-end
-
 # SymmetrySectors interface
 function SymmetrySectors.:×(f1::SectorFusionTree, f2::SectorFusionTree)
   @assert arrows(f1) == arrows(f2)
@@ -154,6 +145,23 @@ end
 # TBD change type depending on AbelianStyle?
 fusiontree_eltype(::Type{<:AbstractSector}) = Float64
 
+# constructors
+function build_trees(legs::Vararg{AbstractGradedUnitRange})
+  # construct all authorized trees for each outer block in legs
+  tree_arrows = isdual.(legs)
+  return mapreduce(vcat, Iterators.product(blocklabels.(legs)...)) do it
+    return build_trees(it, tree_arrows)
+  end
+end
+
+function build_trees(
+  sectors_to_fuse::NTuple{N,<:AbstractSector}, arrows_to_fuse::NTuple{N,Bool}
+) where {N}
+  # construct all authorized trees with fixed outer sectors
+  trees = [SectorFusionTree(first(sectors_to_fuse), first(arrows_to_fuse))]
+  return recursive_build_trees(trees, Base.tail(sectors_to_fuse), Base.tail(arrows_to_fuse))
+end
+
 #
 # =====================================  Internals  ========================================
 #
@@ -205,7 +213,7 @@ function braid_tuples(t1::Tuple{Vararg{Any,N}}, t2::Tuple{Vararg{Any,N}}) where
   return flatten_tuples(nested)
 end
 
-function grow_tree(
+function append_tree_leave(
   parent_tree::SectorFusionTree,
   branch_sector::AbstractSector,
   level_arrow::Bool,
@@ -221,38 +229,33 @@ function grow_tree(
   )
 end
 
-function grow_tree(
+function fuse_next_sector(
   parent_tree::SectorFusionTree, branch_sector::AbstractSector, level_arrow::Bool
 )
   new_space = fusion_product(root_sector(parent_tree), branch_sector)
   return mapreduce(vcat, zip(blocklabels(new_space), blocklengths(new_space))) do (la, n)
     return [
-      grow_tree(parent_tree, branch_sector, level_arrow, la, outer_mult) for
+      append_tree_leave(parent_tree, branch_sector, level_arrow, la, outer_mult) for
       outer_mult in 1:n
     ]
   end
 end
 
-function build_trees(old_trees::Vector, sectors_to_fuse::Tuple, arrows_to_fuse::Tuple)
+function recursive_build_trees(
+  old_trees::Vector, sectors_to_fuse::Tuple, arrows_to_fuse::Tuple
+)
   next_level_trees = mapreduce(vcat, old_trees) do tree
-    return grow_tree(tree, first(sectors_to_fuse), first(arrows_to_fuse))
+    return fuse_next_sector(tree, first(sectors_to_fuse), first(arrows_to_fuse))
   end
-  return build_trees(
+  return recursive_build_trees(
     next_level_trees, Base.tail(sectors_to_fuse), Base.tail(arrows_to_fuse)
   )
 end
 
-function build_trees(trees::Vector, ::Tuple{}, ::Tuple{})
+function recursive_build_trees(trees::Vector, ::Tuple{}, ::Tuple{})
   return trees
 end
 
-function build_trees(
-  sectors_to_fuse::NTuple{N,<:AbstractSector}, arrows_to_fuse::NTuple{N,Bool}
-) where {N}
-  trees = [SectorFusionTree(first(sectors_to_fuse), first(arrows_to_fuse))]
-  return build_trees(trees, Base.tail(sectors_to_fuse), Base.tail(arrows_to_fuse))
-end
-
 # --------------- convert to Array  ---------------
 to_array(::SectorFusionTree{<:Any,0}) = ones(1)
 
diff --git a/src/fusiontensor/fusiontensor.jl b/src/fusiontensor/fusiontensor.jl
index 494bffd..b129f75 100644
--- a/src/fusiontensor/fusiontensor.jl
+++ b/src/fusiontensor/fusiontensor.jl
@@ -165,22 +165,20 @@ end
 function fusion_trees_external_multiplicities(
   outer_legs::Tuple{Vararg{AbstractGradedUnitRange}}
 )
-  return mapreduce(vcat, CartesianIndices(blocklength.(outer_legs))) do it
-    return fusion_trees_external_multiplicities(outer_legs, Tuple(it))
-  end
+  return Iterators.flatten(
+    block_fusion_trees_external_multiplicities.(Iterators.product(blocks.(outer_legs)...))
+  )
 end
 
-function fusion_trees_external_multiplicities(
-  outer_legs::NTuple{N,AbstractGradedUnitRange}, indices::NTuple{N,Int}
+function block_fusion_trees_external_multiplicities(
+  it::NTuple{N,AbstractUnitRange}
 ) where {N}
-  block_sectors = map((g, i) -> blocklabels(g)[i], outer_legs, indices)
-  block_mult = mapreduce((g, i) -> blocklengths(g)[i], *, outer_legs, indices; init=1)
-  return build_trees(block_sectors, isdual.(outer_legs)) .=> block_mult
+  block_sectors = only.(blocklabels.(it))
+  block_mult = prod(length.(it))
+  return build_trees(block_sectors, isdual.(it)) .=> block_mult
 end
 
-function compute_inner_ranges(
-  fusion_trees_mult::AbstractVector{<:Pair{<:SectorFusionTree,<:Integer}}
-)
+function compute_inner_ranges(fusion_trees_mult)
   fused_leg = blockmergesort(
     gradedrange(root_sector.(first.(fusion_trees_mult)) .=> last.(fusion_trees_mult))
   )