-- |
-- 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 (..),
    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)

-- | 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

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