-- |
-- Module      : Tile.Tiling
-- Description : Recursive decompositions of affine tiles.
--
-- A tiling defines the structural decomposition of a tile. The tree
-- layer derives communication children by contracting anchor edges.
module Tile.Tiling
  ( -- * Tiling
    Tiling (..),

    -- * Built-in tilings
    BlockPartitioning (..),
    BoundedFanout (..),
    minimumFanout,
    effectiveFanout,
    Bisection (..),

    -- * Decomposition nodes
    Split (..),
    Relation (..),
    TileNode (..),
  )
where

import Tile.Affine
import Tile.Tile

-- | A recursive decomposition of an affine 'Tile'.
--
-- A 'Tiling' defines the structural children of a tile. Each
-- structural child is labelled by its relationship to the parent: an
-- 'Anchor' child contains the parent root, while a 'Sibling' child
-- introduces a distinct communication root.
--
-- A lawful 'Tiling' satisfies:
--
-- [Inclusion]
--   Every structural child is contained in its parent.
--
--   @
--   tileRanks child \`isSubsetOf\` tileRanks parent
--   @
--
-- [Structural cover]
--   For a non-terminal parent, the structural children partition the
--   parent ranks: their ranks are pairwise disjoint and their union
--   is the parent.
--
-- [Progress]
--   Every structural child is strictly smaller than its parent.
--
--   @
--   length (tileRanks child) < length (tileRanks parent)
--   @
--
-- [Anchor preservation]
--   An anchor child preserves the parent root.
--
--   @
--   case relation node of
--     Anchor _ -> root (tile node) == root parent
--     _        -> True
--   @
--
-- [Sibling movement]
--   A sibling child has a distinct root from its parent.
--
--   @
--   case relation node of
--     Sibling _ -> root (tile node) /= root parent
--     _         -> True
--   @
class Tiling t where
  -- | Structural children of a tile, labelled by their decomposition
  -- relation. Anchor children carry the parent root; sibling children
  -- introduce a distinct communication root.
  childNodes :: t -> Tile -> [TileNode]

-- | The affine dimension and starting index selected by a tiling
-- step.
--
-- A 'Split' records where a child tile came from inside its parent.
data Split = Split
  { -- | Dimension split by the tiling step.
    Split -> Int
dim :: Int,
    -- | Starting index of the child along 'dim'.
    Split -> Int
index :: Int
  }
  deriving (Int -> Split -> ShowS
[Split] -> ShowS
Split -> String
(Int -> Split -> ShowS)
-> (Split -> String) -> ([Split] -> ShowS) -> Show Split
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Split -> ShowS
showsPrec :: Int -> Split -> ShowS
$cshow :: Split -> String
show :: Split -> String
$cshowList :: [Split] -> ShowS
showList :: [Split] -> ShowS
Show, Split -> Split -> Bool
(Split -> Split -> Bool) -> (Split -> Split -> Bool) -> Eq Split
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Split -> Split -> Bool
== :: Split -> Split -> Bool
$c/= :: Split -> Split -> Bool
/= :: Split -> Split -> Bool
Eq, Eq Split
Eq Split =>
(Split -> Split -> Ordering)
-> (Split -> Split -> Bool)
-> (Split -> Split -> Bool)
-> (Split -> Split -> Bool)
-> (Split -> Split -> Bool)
-> (Split -> Split -> Split)
-> (Split -> Split -> Split)
-> Ord Split
Split -> Split -> Bool
Split -> Split -> Ordering
Split -> Split -> Split
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Split -> Split -> Ordering
compare :: Split -> Split -> Ordering
$c< :: Split -> Split -> Bool
< :: Split -> Split -> Bool
$c<= :: Split -> Split -> Bool
<= :: Split -> Split -> Bool
$c> :: Split -> Split -> Bool
> :: Split -> Split -> Bool
$c>= :: Split -> Split -> Bool
>= :: Split -> Split -> Bool
$cmax :: Split -> Split -> Split
max :: Split -> Split -> Split
$cmin :: Split -> Split -> Split
min :: Split -> Split -> Split
Ord)

-- | Relationship between a 'TileNode' and its parent.
--
-- Relations distinguish geometry-only anchor steps from communication
-- steps. The tree layer contracts anchor edges; sibling edges become
-- communication edges.
data Relation
  = -- | The root of a decomposition tree.
    Root
  | -- | A child that preserves the parent root.
    Anchor Split
  | -- | A child with a distinct root from the parent.
    Sibling Split
  deriving (Int -> Relation -> ShowS
[Relation] -> ShowS
Relation -> String
(Int -> Relation -> ShowS)
-> (Relation -> String) -> ([Relation] -> ShowS) -> Show Relation
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Relation -> ShowS
showsPrec :: Int -> Relation -> ShowS
$cshow :: Relation -> String
show :: Relation -> String
$cshowList :: [Relation] -> ShowS
showList :: [Relation] -> ShowS
Show, Relation -> Relation -> Bool
(Relation -> Relation -> Bool)
-> (Relation -> Relation -> Bool) -> Eq Relation
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Relation -> Relation -> Bool
== :: Relation -> Relation -> Bool
$c/= :: Relation -> Relation -> Bool
/= :: Relation -> Relation -> Bool
Eq, Eq Relation
Eq Relation =>
(Relation -> Relation -> Ordering)
-> (Relation -> Relation -> Bool)
-> (Relation -> Relation -> Bool)
-> (Relation -> Relation -> Bool)
-> (Relation -> Relation -> Bool)
-> (Relation -> Relation -> Relation)
-> (Relation -> Relation -> Relation)
-> Ord Relation
Relation -> Relation -> Bool
Relation -> Relation -> Ordering
Relation -> Relation -> Relation
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Relation -> Relation -> Ordering
compare :: Relation -> Relation -> Ordering
$c< :: Relation -> Relation -> Bool
< :: Relation -> Relation -> Bool
$c<= :: Relation -> Relation -> Bool
<= :: Relation -> Relation -> Bool
$c> :: Relation -> Relation -> Bool
> :: Relation -> Relation -> Bool
$c>= :: Relation -> Relation -> Bool
>= :: Relation -> Relation -> Bool
$cmax :: Relation -> Relation -> Relation
max :: Relation -> Relation -> Relation
$cmin :: Relation -> Relation -> Relation
min :: Relation -> Relation -> Relation
Ord)

-- | A tile labelled with its relationship to a parent tile.
--
-- 'TileNode' is the node type used by decomposition and hop trees.
-- The 'relation' field records how the node arises from its parent;
-- the root node of a tree uses 'Root'.
data TileNode = TileNode
  { -- | Tile carried by the node.
    TileNode -> Tile
tile :: Tile,
    -- | Relationship to the parent node.
    TileNode -> Relation
relation :: Relation
  }
  deriving (Int -> TileNode -> ShowS
[TileNode] -> ShowS
TileNode -> String
(Int -> TileNode -> ShowS)
-> (TileNode -> String) -> ([TileNode] -> ShowS) -> Show TileNode
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> TileNode -> ShowS
showsPrec :: Int -> TileNode -> ShowS
$cshow :: TileNode -> String
show :: TileNode -> String
$cshowList :: [TileNode] -> ShowS
showList :: [TileNode] -> ShowS
Show, TileNode -> TileNode -> Bool
(TileNode -> TileNode -> Bool)
-> (TileNode -> TileNode -> Bool) -> Eq TileNode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: TileNode -> TileNode -> Bool
== :: TileNode -> TileNode -> Bool
$c/= :: TileNode -> TileNode -> Bool
/= :: TileNode -> TileNode -> Bool
Eq)

-- | Minimum communication fan-out required by rectangular geometry.
--
-- This is the number of active dimensions in the tile. With affine
-- rectangular children and a corner root, each active dimension
-- contributes one necessary frontier region away from the root.
minimumFanout :: Tile -> Int
minimumFanout :: Tile -> Int
minimumFanout Tile
tile =
  [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int
n | Int
n <- AffineRankSpace -> [Int]
sizes (Tile -> AffineRankSpace
space Tile
tile), Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1]

-- | Fan-out actually available to 'BoundedFanout'.
--
-- The requested cap is honored when geometry permits it. If the cap
-- is below the rectangular minimum, the geometric minimum is used
-- instead.
effectiveFanout :: Tile -> Int -> Int
effectiveFanout :: Tile -> Int -> Int
effectiveFanout Tile
tile Int
requested =
  Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
requested (Tile -> Int
minimumFanout Tile
tile)

-- | Fix one affine dimension of a tile to a single index.
fixTileDim :: Tile -> Int -> Int -> Maybe Tile
fixTileDim :: Tile -> Int -> Int -> Maybe Tile
fixTileDim Tile
tile Int
dim Int
i = AffineRankSpace -> Tile
Tile (AffineRankSpace -> Tile) -> Maybe AffineRankSpace -> Maybe Tile
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AffineRankSpace -> Int -> Int -> Maybe AffineRankSpace
fixDim (Tile -> AffineRankSpace
space Tile
tile) Int
dim Int
i

-- | Select a contiguous interval along one tile dimension.
selectTileDim :: Tile -> Int -> Int -> Int -> Maybe Tile
selectTileDim :: Tile -> Int -> Int -> Int -> Maybe Tile
selectTileDim Tile
tile Int
dim Int
begin Int
end =
  AffineRankSpace -> Tile
Tile (AffineRankSpace -> Tile) -> Maybe AffineRankSpace -> Maybe Tile
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> AffineRankSpace
-> Int -> Int -> Int -> Int -> Maybe AffineRankSpace
select (Tile -> AffineRankSpace
space Tile
tile) Int
dim Int
begin Int
end Int
1

-- | Partition by fixing one coordinate at a time.
--
-- 'BlockPartitioning' finds the first non-singleton dimension. It
-- creates one anchor child at index @0@ and one sibling child for
-- each remaining index in that dimension.
data BlockPartitioning = BlockPartitioning
  deriving (Int -> BlockPartitioning -> ShowS
[BlockPartitioning] -> ShowS
BlockPartitioning -> String
(Int -> BlockPartitioning -> ShowS)
-> (BlockPartitioning -> String)
-> ([BlockPartitioning] -> ShowS)
-> Show BlockPartitioning
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BlockPartitioning -> ShowS
showsPrec :: Int -> BlockPartitioning -> ShowS
$cshow :: BlockPartitioning -> String
show :: BlockPartitioning -> String
$cshowList :: [BlockPartitioning] -> ShowS
showList :: [BlockPartitioning] -> ShowS
Show, BlockPartitioning -> BlockPartitioning -> Bool
(BlockPartitioning -> BlockPartitioning -> Bool)
-> (BlockPartitioning -> BlockPartitioning -> Bool)
-> Eq BlockPartitioning
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BlockPartitioning -> BlockPartitioning -> Bool
== :: BlockPartitioning -> BlockPartitioning -> Bool
$c/= :: BlockPartitioning -> BlockPartitioning -> Bool
/= :: BlockPartitioning -> BlockPartitioning -> Bool
Eq)

instance Tiling BlockPartitioning where
  childNodes :: BlockPartitioning -> Tile -> [TileNode]
childNodes BlockPartitioning
_ Tile
tile = Tile -> Int -> [TileNode]
go Tile
tile Int
0
    where
      go :: Tile -> Int -> [TileNode]
go Tile
t Int
d
        | Int
d Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (AffineRankSpace -> [Int]
sizes (Tile -> AffineRankSpace
space Tile
t)) = []
        | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = Tile -> Int -> [TileNode]
go Tile
t (Int
d Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
        | Bool
otherwise =
            let siblings :: [TileNode]
siblings =
                  [ Tile -> Relation -> TileNode
TileNode Tile
child (Split -> Relation
Sibling (Int -> Int -> Split
Split Int
d Int
i))
                  | Int
i <- [Int
1 .. Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1],
                    Just Tile
child <- [Tile -> Int -> Int -> Maybe Tile
fixTileDim Tile
t Int
d Int
i]
                  ]
                anchors :: [TileNode]
anchors =
                  [ Tile -> Relation -> TileNode
TileNode Tile
child (Split -> Relation
Anchor (Int -> Int -> Split
Split Int
d Int
0))
                  | Just Tile
child <- [Tile -> Int -> Int -> Maybe Tile
fixTileDim Tile
t Int
d Int
0]
                  ]
             in [TileNode]
siblings [TileNode] -> [TileNode] -> [TileNode]
forall a. [a] -> [a] -> [a]
++ [TileNode]
anchors
        where
          n :: Int
n = AffineRankSpace -> [Int]
sizes (Tile -> AffineRankSpace
space Tile
t) [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
d

-- | Partition by bisecting one dimension at a time.
--
-- 'Bisection' finds the first non-singleton dimension. It keeps the
-- lower @floor(n / 2)@ half as the anchor and creates one sibling
-- from the remaining upper half.
data Bisection = Bisection
  deriving (Int -> Bisection -> ShowS
[Bisection] -> ShowS
Bisection -> String
(Int -> Bisection -> ShowS)
-> (Bisection -> String)
-> ([Bisection] -> ShowS)
-> Show Bisection
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Bisection -> ShowS
showsPrec :: Int -> Bisection -> ShowS
$cshow :: Bisection -> String
show :: Bisection -> String
$cshowList :: [Bisection] -> ShowS
showList :: [Bisection] -> ShowS
Show, Bisection -> Bisection -> Bool
(Bisection -> Bisection -> Bool)
-> (Bisection -> Bisection -> Bool) -> Eq Bisection
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Bisection -> Bisection -> Bool
== :: Bisection -> Bisection -> Bool
$c/= :: Bisection -> Bisection -> Bool
/= :: Bisection -> Bisection -> Bool
Eq)

instance Tiling Bisection where
  childNodes :: Bisection -> Tile -> [TileNode]
childNodes Bisection
_ Tile
tile =
    case Tile -> Maybe Int
firstNonSingletonDim Tile
tile of
      Maybe Int
Nothing -> []
      Just Int
d ->
        let n :: Int
n = AffineRankSpace -> [Int]
sizes (Tile -> AffineRankSpace
space Tile
tile) [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
d
            lower :: Int
lower = Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2
            siblings :: [TileNode]
siblings =
              [ Tile -> Relation -> TileNode
TileNode Tile
child (Split -> Relation
Sibling (Int -> Int -> Split
Split Int
d Int
lower))
              | Just Tile
child <- [Tile -> Int -> Int -> Int -> Maybe Tile
selectTileDim Tile
tile Int
d Int
lower Int
n]
              ]
            anchors :: [TileNode]
anchors =
              [ Tile -> Relation -> TileNode
TileNode Tile
child (Split -> Relation
Anchor (Int -> Int -> Split
Split Int
d Int
0))
              | Just Tile
child <- [Tile -> Int -> Int -> Int -> Maybe Tile
selectTileDim Tile
tile Int
d Int
0 Int
lower]
              ]
         in [TileNode]
siblings [TileNode] -> [TileNode] -> [TileNode]
forall a. [a] -> [a] -> [a]
++ [TileNode]
anchors

-- | Rectangular tiling with bounded communication fan-out.
--
-- 'BoundedFanout' computes the full local rectangular frontier of a
-- tile. The frontier contains, for each active dimension, the ranks
-- whose first coordinate away from the root occurs in that dimension.
-- The remaining root point is emitted as a terminal anchor.
--
-- The requested cap is respected when it is geometrically feasible:
--
-- @
-- minimumFanout tile <= fanout
--   ==> length (children (BoundedFanout fanout) tile) <= fanout
-- @
--
-- Unconditionally:
--
-- @
-- length (children (BoundedFanout fanout) tile)
--   <= effectiveFanout tile fanout
-- @
--
-- The tiler preserves affine rectangular children; it does not
-- introduce jagged regions. A non-positive requested fan-out
-- produces no children.
newtype BoundedFanout = BoundedFanout
  { -- | Requested communication fan-out cap.
    BoundedFanout -> Int
fanout :: Int
  }
  deriving (Int -> BoundedFanout -> ShowS
[BoundedFanout] -> ShowS
BoundedFanout -> String
(Int -> BoundedFanout -> ShowS)
-> (BoundedFanout -> String)
-> ([BoundedFanout] -> ShowS)
-> Show BoundedFanout
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BoundedFanout -> ShowS
showsPrec :: Int -> BoundedFanout -> ShowS
$cshow :: BoundedFanout -> String
show :: BoundedFanout -> String
$cshowList :: [BoundedFanout] -> ShowS
showList :: [BoundedFanout] -> ShowS
Show, BoundedFanout -> BoundedFanout -> Bool
(BoundedFanout -> BoundedFanout -> Bool)
-> (BoundedFanout -> BoundedFanout -> Bool) -> Eq BoundedFanout
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BoundedFanout -> BoundedFanout -> Bool
== :: BoundedFanout -> BoundedFanout -> Bool
$c/= :: BoundedFanout -> BoundedFanout -> Bool
/= :: BoundedFanout -> BoundedFanout -> Bool
Eq)

instance Tiling BoundedFanout where
  childNodes :: BoundedFanout -> Tile -> [TileNode]
childNodes (BoundedFanout Int
requestedFanout) Tile
baseTile
    | Int
requestedFanout Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = []
    | [(Int, Int)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Int, Int)]
activeDims = []
    | Bool
otherwise = [TileNode]
siblings [TileNode] -> [TileNode] -> [TileNode]
forall a. [a] -> [a] -> [a]
++ [TileNode]
anchors
    where
      activeDims :: [(Int, Int)]
activeDims =
        [ (Int
d, Int
n)
        | (Int
d, Int
n) <- [Int] -> [Int] -> [(Int, Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int
0 ..] (AffineRankSpace -> [Int]
sizes (Tile -> AffineRankSpace
space Tile
baseTile)),
          Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
1
        ]

      groupCounts :: [Int]
groupCounts =
        Int -> [(Int, Int)] -> [Int]
allocateGroups (Tile -> Int -> Int
effectiveFanout Tile
baseTile Int
requestedFanout) [(Int, Int)]
activeDims

      siblings :: [TileNode]
siblings =
        [[TileNode]] -> [TileNode]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
          [ Int -> Int -> [TileNode]
frontierChildren Int
d Int
groupCount
          | ((Int
d, Int
_), Int
groupCount) <- [(Int, Int)] -> [Int] -> [((Int, Int), Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip [(Int, Int)]
activeDims [Int]
groupCounts
          ]

      anchors :: [TileNode]
anchors =
        [ Tile -> Relation -> TileNode
TileNode Tile
child (Split -> Relation
Anchor (Int -> Int -> Split
Split Int
d Int
0))
        | -- The terminal root-point anchor belongs to every active
          -- dimension; use the last dimension as its structural label.
          let d :: Int
d = (Int, Int) -> Int
forall a b. (a, b) -> a
fst ([(Int, Int)] -> (Int, Int)
forall a. HasCallStack => [a] -> a
last [(Int, Int)]
activeDims),
          Just Tile
child <- [Tile -> [(Int, Int)] -> Maybe Tile
rootPointTile Tile
baseTile [(Int, Int)]
activeDims]
        ]

      frontierChildren :: Int -> Int -> [TileNode]
frontierChildren Int
d Int
groupCount =
        [ Tile -> Relation -> TileNode
TileNode Tile
child (Split -> Relation
Sibling (Int -> Int -> Split
Split Int
d Int
begin))
        | (Int
begin, Int
end) <- Int -> Int -> [(Int, Int)]
boundedIntervals Int
groupCount (AffineRankSpace -> [Int]
sizes (Tile -> AffineRankSpace
space Tile
baseTile) [Int] -> Int -> Int
forall a. HasCallStack => [a] -> Int -> a
!! Int
d),
          Just Tile
child <- [Tile -> [(Int, Int)] -> Int -> Int -> Int -> Maybe Tile
frontierTile Tile
baseTile [(Int, Int)]
activeDims Int
d Int
begin Int
end]
        ]

allocateGroups :: Int -> [(Int, Int)] -> [Int]
allocateGroups :: Int -> [(Int, Int)] -> [Int]
allocateGroups Int
available [(Int, Int)]
activeDims =
  Int -> [Int] -> [Int] -> [Int]
forall {t}. (Ord t, Num t) => t -> [t] -> [t] -> [t]
go Int
extra [Int]
base [Int]
capacities
  where
    base :: [Int]
base = Int -> Int -> [Int]
forall a. Int -> a -> [a]
replicate ([(Int, Int)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Int)]
activeDims) Int
1
    capacities :: [Int]
capacities = [Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 | (Int
_, Int
n) <- [(Int, Int)]
activeDims]
    extra :: Int
extra = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
0 (Int
available Int -> Int -> Int
forall a. Num a => a -> a -> a
- [(Int, Int)] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [(Int, Int)]
activeDims)

    go :: t -> [t] -> [t] -> [t]
go t
_ [] [] = []
    go t
remaining (t
count : [t]
counts) (t
capacity : [t]
rest) =
      let additional :: t
additional = t -> t -> t
forall a. Ord a => a -> a -> a
min t
remaining (t
capacity t -> t -> t
forall a. Num a => a -> a -> a
- t
count)
       in t
count t -> t -> t
forall a. Num a => a -> a -> a
+ t
additional t -> [t] -> [t]
forall a. a -> [a] -> [a]
: t -> [t] -> [t] -> [t]
go (t
remaining t -> t -> t
forall a. Num a => a -> a -> a
- t
additional) [t]
counts [t]
rest
    go t
_ [t]
_ [t]
_ = String -> [t]
forall a. HasCallStack => String -> a
error String
"allocateGroups: mismatched group and capacity lists"

boundedIntervals :: Int -> Int -> [(Int, Int)]
boundedIntervals :: Int -> Int -> [(Int, Int)]
boundedIntervals Int
groupCount Int
extent
  | Int
groupCount Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = []
  | Int
remaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = []
  | Bool
otherwise = Int -> [Int] -> [(Int, Int)]
forall {t}. Num t => t -> [t] -> [(t, t)]
go Int
1 [Int]
groupSizes
  where
    remaining :: Int
remaining = Int
extent Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
    groups :: Int
groups = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
groupCount Int
remaining
    (Int
base, Int
extra) = Int
remaining Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`divMod` Int
groups

    groupSizes :: [Int]
groupSizes =
      [ Int
base Int -> Int -> Int
forall a. Num a => a -> a -> a
+ if Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
extra then Int
1 else Int
0
      | Int
i <- [Int
0 .. Int
groups Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1]
      ]

    go :: t -> [t] -> [(t, t)]
go t
_ [] = []
    go t
begin (t
groupSize : [t]
rest) =
      let end :: t
end = t
begin t -> t -> t
forall a. Num a => a -> a -> a
+ t
groupSize
       in (t
begin, t
end) (t, t) -> [(t, t)] -> [(t, t)]
forall a. a -> [a] -> [a]
: t -> [t] -> [(t, t)]
go t
end [t]
rest

frontierTile :: Tile -> [(Int, Int)] -> Int -> Int -> Int -> Maybe Tile
frontierTile :: Tile -> [(Int, Int)] -> Int -> Int -> Int -> Maybe Tile
frontierTile Tile
baseTile [(Int, Int)]
activeDims Int
dim Int
begin Int
end = do
  anchored <- Tile -> [(Int, Int)] -> Int -> Maybe Tile
anchorPrefix Tile
baseTile [(Int, Int)]
activeDims Int
dim
  selectTileDim anchored dim begin end

anchorPrefix :: Tile -> [(Int, Int)] -> Int -> Maybe Tile
anchorPrefix :: Tile -> [(Int, Int)] -> Int -> Maybe Tile
anchorPrefix Tile
baseTile [(Int, Int)]
activeDims Int
dim =
  (Maybe Tile -> (Int, Int) -> Maybe Tile)
-> Maybe Tile -> [(Int, Int)] -> Maybe Tile
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
    (\Maybe Tile
mt (Int
d, Int
_) -> Maybe Tile
mt Maybe Tile -> (Tile -> Maybe Tile) -> Maybe Tile
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Tile
tile -> Tile -> Int -> Int -> Maybe Tile
fixTileDim Tile
tile Int
d Int
0)
    (Tile -> Maybe Tile
forall a. a -> Maybe a
Just Tile
baseTile)
    (((Int, Int) -> Bool) -> [(Int, Int)] -> [(Int, Int)]
forall a. (a -> Bool) -> [a] -> [a]
takeWhile ((Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
dim) (Int -> Bool) -> ((Int, Int) -> Int) -> (Int, Int) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Int, Int) -> Int
forall a b. (a, b) -> a
fst) [(Int, Int)]
activeDims)

rootPointTile :: Tile -> [(Int, Int)] -> Maybe Tile
rootPointTile :: Tile -> [(Int, Int)] -> Maybe Tile
rootPointTile =
  (Maybe Tile -> (Int, Int) -> Maybe Tile)
-> Maybe Tile -> [(Int, Int)] -> Maybe Tile
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl
    ( \Maybe Tile
mt (Int
d, Int
_) ->
        Maybe Tile
mt Maybe Tile -> (Tile -> Maybe Tile) -> Maybe Tile
forall a b. Maybe a -> (a -> Maybe b) -> Maybe b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \Tile
tile -> Tile -> Int -> Int -> Maybe Tile
fixTileDim Tile
tile Int
d Int
0
    )
    (Maybe Tile -> [(Int, Int)] -> Maybe Tile)
-> (Tile -> Maybe Tile) -> Tile -> [(Int, Int)] -> Maybe Tile
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tile -> Maybe Tile
forall a. a -> Maybe a
Just

firstNonSingletonDim :: Tile -> Maybe Int
firstNonSingletonDim :: Tile -> Maybe Int
firstNonSingletonDim Tile
tile =
  Int -> [Int] -> Maybe Int
forall {a} {t}. (Ord a, Num a, Num t) => t -> [a] -> Maybe t
go Int
0 (AffineRankSpace -> [Int]
sizes (Tile -> AffineRankSpace
space Tile
tile))
  where
    go :: t -> [a] -> Maybe t
go t
_ [] = Maybe t
forall a. Maybe a
Nothing
    go t
d (a
n : [a]
ns)
      | a
n a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
1 = t -> Maybe t
forall a. a -> Maybe a
Just t
d
      | Bool
otherwise = t -> [a] -> Maybe t
go (t
d t -> t -> t
forall a. Num a => a -> a -> a
+ t
1) [a]
ns