refactor repacking of transformer structure

Author: lkdvos

src/tensors/treetransformers.jl

@@ -59,7 +59,9 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc)
     t₀ = Base.time()
     permute(Vsrc, p) == Vdst || throw(SpaceMismatch("Incompatible spaces for permuting."))
     structure_dst = fusionblockstructure(Vdst)
+    fusionstructure_dst = structure_dst.fusiontreestructure
     structure_src = fusionblockstructure(Vsrc)
+    fusionstructure_src = structure_src.fusiontreestructure
     I = sectortype(Vsrc)
 
     uncoupleds_src = map(structure_src.fusiontreelist) do (f₁, f₂)
@@ -94,13 +96,10 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc)
             end
         end
 
-        structs_src = structure_src.fusiontreestructure[ids_src]
-        sz_src = structs_src[1][1]
-        newstructs_src = map(x -> (x[2], x[3]), structs_src)
-
-        structs_dst = structure_dst.fusiontreestructure[ids_dst]
-        sz_dst = structs_dst[1][1]
-        newstructs_dst = map(x -> (x[2], x[3]), structs_dst)
+        # size is shared between blocks, so repack:
+        # from [(sz, strides, offset), ...] to (sz, [(strides, offset), ...])
+        sz_src, newstructs_src = repack_transformer_structure(fusionstructure_src, ids_src)
+        sz_dst, newstructs_dst = repack_transformer_structure(fusionstructure_dst, ids_dst)
 
         @debug("Created recoupling block for uncoupled: $uncoupled",
                sz = size(matrix), sparsity = count(!iszero, matrix) / length(matrix))
@@ -124,6 +123,12 @@ function GenericTreeTransformer(transform, p, Vdst, Vsrc)
     return transformer
 end
 
+function repack_transformer_structure(structures, ids)
+    sz = structures[first(ids)][1]
+    strides_offsets = map(i -> (structures[i][2], structures[i][3]), ids)
+    return sz, strides_offsets
+end
+
 function buffersize(transformer::GenericTreeTransformer)
     return maximum(transformer.data; init=0) do (basistransform, structures_dst, _)
         return prod(structures_dst[1]) * size(basistransform, 1)