def repartition

stelliferous@0.0.2.0 · tensor_band.bend:325

import Base
import stelliferous@0.0.2.0/storage_buffer.bend as Storage
import stelliferous@0.0.2.0/traversal_partition.bend as Partition

325def repartition(~Element: Data,+geometry: Geometry,+fill: Element,layout: Plan,bands: Bands(~Element)) -> Bands(~Element) & Bands(~Element):
326  match layout:
327    case Rows{+first,+rows}:
328      Storage.bind_pair(~Bands(~Element),~Storage.Buffer<Element>,~(Bands(~Element) & Bands(~Element)),
329        range(~Element,geometry,first,rows,bands,allocate(~Element,fill,local_count(geometry,rows))),source => output =>
330          (source,Partition.Leaf{output,first,rows,Unit{}}))
331    case Split{left,right}:
332      Storage.bind_pair(~Bands(~Element),~Bands(~Element),~(Bands(~Element) & Bands(~Element)),
333        repartition(~Element,geometry,fill,left,bands),source => left =>
334          Storage.bind_pair(~Bands(~Element),~Bands(~Element),~(Bands(~Element) & Bands(~Element)),
335            repartition(~Element,geometry,fill,right,source),source => right => (source,Partition.Fork{left,right})))