module Tile.Tree
  ( Tree (..),
    DecompositionView,
    HopView,
    TileTree (..),
    DecompositionTree,
    HopTree,
    mapTree,
    unfoldTree,
    renderTreeWith,
    SendTree (..),
    RoutedTree (..),
    scheduleTree,
    routedTree,
    renderRoutedTree,
    decompositionTree,
    contractAnchors,
    hopTree,
    sendTree,
    children,
    renderDecompositionTree,
    renderHopTree,
    renderSendTree,
  )
where

import Data.List (sortOn)
import Data.Map.Strict qualified as Map
import Tile.Geometry
import Tile.Schedule
import Tile.Tile
import Tile.Tiling

data Tree a = Tree
  { forall a. Tree a -> a
treeLabel :: a,
    forall a. Tree a -> [Tree a]
subtrees :: [Tree a]
  }
  deriving (Int -> Tree a -> ShowS
[Tree a] -> ShowS
Tree a -> String
(Int -> Tree a -> ShowS)
-> (Tree a -> String) -> ([Tree a] -> ShowS) -> Show (Tree a)
forall a. Show a => Int -> Tree a -> ShowS
forall a. Show a => [Tree a] -> ShowS
forall a. Show a => Tree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Tree a -> ShowS
showsPrec :: Int -> Tree a -> ShowS
$cshow :: forall a. Show a => Tree a -> String
show :: Tree a -> String
$cshowList :: forall a. Show a => [Tree a] -> ShowS
showList :: [Tree a] -> ShowS
Show, Tree a -> Tree a -> Bool
(Tree a -> Tree a -> Bool)
-> (Tree a -> Tree a -> Bool) -> Eq (Tree a)
forall a. Eq a => Tree a -> Tree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => Tree a -> Tree a -> Bool
== :: Tree a -> Tree a -> Bool
$c/= :: forall a. Eq a => Tree a -> Tree a -> Bool
/= :: Tree a -> Tree a -> Bool
Eq)

data DecompositionView

data HopView

newtype TileTree view = TileTree
  { forall {k} (view :: k). TileTree view -> Tree TileNode
getTileTree :: Tree TileNode
  }
  deriving (Int -> TileTree view -> ShowS
[TileTree view] -> ShowS
TileTree view -> String
(Int -> TileTree view -> ShowS)
-> (TileTree view -> String)
-> ([TileTree view] -> ShowS)
-> Show (TileTree view)
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
forall k (view :: k). Int -> TileTree view -> ShowS
forall k (view :: k). [TileTree view] -> ShowS
forall k (view :: k). TileTree view -> String
$cshowsPrec :: forall k (view :: k). Int -> TileTree view -> ShowS
showsPrec :: Int -> TileTree view -> ShowS
$cshow :: forall k (view :: k). TileTree view -> String
show :: TileTree view -> String
$cshowList :: forall k (view :: k). [TileTree view] -> ShowS
showList :: [TileTree view] -> ShowS
Show, TileTree view -> TileTree view -> Bool
(TileTree view -> TileTree view -> Bool)
-> (TileTree view -> TileTree view -> Bool) -> Eq (TileTree view)
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
forall k (view :: k). TileTree view -> TileTree view -> Bool
$c== :: forall k (view :: k). TileTree view -> TileTree view -> Bool
== :: TileTree view -> TileTree view -> Bool
$c/= :: forall k (view :: k). TileTree view -> TileTree view -> Bool
/= :: TileTree view -> TileTree view -> Bool
Eq)

type DecompositionTree = TileTree DecompositionView

type HopTree = TileTree HopView

newtype SendTree = SendTree
  { SendTree -> Tree Tile
getSendTree :: Tree Tile
  }
  deriving (Int -> SendTree -> ShowS
[SendTree] -> ShowS
SendTree -> String
(Int -> SendTree -> ShowS)
-> (SendTree -> String) -> ([SendTree] -> ShowS) -> Show SendTree
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> SendTree -> ShowS
showsPrec :: Int -> SendTree -> ShowS
$cshow :: SendTree -> String
show :: SendTree -> String
$cshowList :: [SendTree] -> ShowS
showList :: [SendTree] -> ShowS
Show, SendTree -> SendTree -> Bool
(SendTree -> SendTree -> Bool)
-> (SendTree -> SendTree -> Bool) -> Eq SendTree
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: SendTree -> SendTree -> Bool
== :: SendTree -> SendTree -> Bool
$c/= :: SendTree -> SendTree -> Bool
/= :: SendTree -> SendTree -> Bool
Eq)

newtype RoutedTree a = RoutedTree
  { forall a. RoutedTree a -> Tree a
getRoutedTree :: Tree a
  }
  deriving (Int -> RoutedTree a -> ShowS
[RoutedTree a] -> ShowS
RoutedTree a -> String
(Int -> RoutedTree a -> ShowS)
-> (RoutedTree a -> String)
-> ([RoutedTree a] -> ShowS)
-> Show (RoutedTree a)
forall a. Show a => Int -> RoutedTree a -> ShowS
forall a. Show a => [RoutedTree a] -> ShowS
forall a. Show a => RoutedTree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> RoutedTree a -> ShowS
showsPrec :: Int -> RoutedTree a -> ShowS
$cshow :: forall a. Show a => RoutedTree a -> String
show :: RoutedTree a -> String
$cshowList :: forall a. Show a => [RoutedTree a] -> ShowS
showList :: [RoutedTree a] -> ShowS
Show, RoutedTree a -> RoutedTree a -> Bool
(RoutedTree a -> RoutedTree a -> Bool)
-> (RoutedTree a -> RoutedTree a -> Bool) -> Eq (RoutedTree a)
forall a. Eq a => RoutedTree a -> RoutedTree a -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: forall a. Eq a => RoutedTree a -> RoutedTree a -> Bool
== :: RoutedTree a -> RoutedTree a -> Bool
$c/= :: forall a. Eq a => RoutedTree a -> RoutedTree a -> Bool
/= :: RoutedTree a -> RoutedTree a -> Bool
Eq)

decompositionTree :: (Tiling t) => t -> Tile -> DecompositionTree
decompositionTree :: forall t. Tiling t => t -> Tile -> DecompositionTree
decompositionTree t
tiling Tile
baseTile =
  Tree TileNode -> DecompositionTree
forall {k} (view :: k). Tree TileNode -> TileTree view
TileTree (Tree TileNode -> DecompositionTree)
-> Tree TileNode -> DecompositionTree
forall a b. (a -> b) -> a -> b
$
    (TileNode -> [TileNode]) -> TileNode -> Tree TileNode
forall a. (a -> [a]) -> a -> Tree a
unfoldTree
      (t -> Tile -> [TileNode]
forall t. Tiling t => t -> Tile -> [TileNode]
childNodes t
tiling (Tile -> [TileNode])
-> (TileNode -> Tile) -> TileNode -> [TileNode]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TileNode -> Tile
tile)
      (Tile -> Relation -> TileNode
TileNode Tile
baseTile Relation
Root)

contractAnchors :: DecompositionTree -> HopTree
contractAnchors :: DecompositionTree -> HopTree
contractAnchors (TileTree Tree TileNode
tree) = Tree TileNode -> HopTree
forall {k} (view :: k). Tree TileNode -> TileTree view
TileTree (Tree TileNode -> Tree TileNode
go Tree TileNode
tree)
  where
    go :: Tree TileNode -> Tree TileNode
go (Tree TileNode
node [Tree TileNode]
kids) = TileNode -> [Tree TileNode] -> Tree TileNode
forall a. a -> [Tree a] -> Tree a
Tree TileNode
node ((Tree TileNode -> [Tree TileNode])
-> [Tree TileNode] -> [Tree TileNode]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree TileNode -> [Tree TileNode]
project [Tree TileNode]
kids)
    project :: Tree TileNode -> [Tree TileNode]
project subtree :: Tree TileNode
subtree@(Tree TileNode
node [Tree TileNode]
kids) =
      case TileNode -> Relation
relation TileNode
node of
        Sibling Split
_ -> [Tree TileNode -> Tree TileNode
go Tree TileNode
subtree]
        Anchor Split
_ -> (Tree TileNode -> [Tree TileNode])
-> [Tree TileNode] -> [Tree TileNode]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Tree TileNode -> [Tree TileNode]
project [Tree TileNode]
kids
        Relation
Root -> []

hopTree :: (Tiling t) => t -> Tile -> HopTree
hopTree :: forall t. Tiling t => t -> Tile -> HopTree
hopTree t
tiling = DecompositionTree -> HopTree
contractAnchors (DecompositionTree -> HopTree)
-> (Tile -> DecompositionTree) -> Tile -> HopTree
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> Tile -> DecompositionTree
forall t. Tiling t => t -> Tile -> DecompositionTree
decompositionTree t
tiling

sendTree :: (Tiling t) => t -> Tile -> SendTree
sendTree :: forall t. Tiling t => t -> Tile -> SendTree
sendTree t
tiling =
  Tree Tile -> SendTree
SendTree (Tree Tile -> SendTree) -> (Tile -> Tree Tile) -> Tile -> SendTree
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (TileNode -> Tile) -> Tree TileNode -> Tree Tile
forall a b. (a -> b) -> Tree a -> Tree b
mapTree TileNode -> Tile
tile (Tree TileNode -> Tree Tile)
-> (Tile -> Tree TileNode) -> Tile -> Tree Tile
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HopTree -> Tree TileNode
forall {k} (view :: k). TileTree view -> Tree TileNode
getTileTree (HopTree -> Tree TileNode)
-> (Tile -> HopTree) -> Tile -> Tree TileNode
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> Tile -> HopTree
forall t. Tiling t => t -> Tile -> HopTree
hopTree t
tiling

children :: (Tiling t) => t -> Tile -> [Tile]
children :: forall t. Tiling t => t -> Tile -> [Tile]
children t
tiling =
  (Tree Tile -> Tile) -> [Tree Tile] -> [Tile]
forall a b. (a -> b) -> [a] -> [b]
map Tree Tile -> Tile
forall a. Tree a -> a
treeLabel ([Tree Tile] -> [Tile]) -> (Tile -> [Tree Tile]) -> Tile -> [Tile]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree Tile -> [Tree Tile]
forall a. Tree a -> [Tree a]
subtrees (Tree Tile -> [Tree Tile])
-> (Tile -> Tree Tile) -> Tile -> [Tree Tile]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SendTree -> Tree Tile
getSendTree (SendTree -> Tree Tile) -> (Tile -> SendTree) -> Tile -> Tree Tile
forall b c a. (b -> c) -> (a -> b) -> a -> c
. t -> Tile -> SendTree
forall t. Tiling t => t -> Tile -> SendTree
sendTree t
tiling

scheduleTree :: (Ord a) => a -> Schedule a -> RoutedTree a
scheduleTree :: forall a. Ord a => a -> Schedule a -> RoutedTree a
scheduleTree a
ingress Schedule a
schedule =
  Tree a -> RoutedTree a
forall a. Tree a -> RoutedTree a
RoutedTree ((a -> [a]) -> a -> Tree a
forall a. (a -> [a]) -> a -> Tree a
unfoldTree a -> [a]
childrenOf a
ingress)
  where
    childrenOf :: a -> [a]
childrenOf a
member =
      [a] -> a -> Map a [a] -> [a]
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault [] a
member (Schedule a -> Map a [a]
forall a. Ord a => Schedule a -> Map a [a]
adjacencyList Schedule a
schedule)

routedTree :: (Ord a) => RoutedSchedule a -> RoutedTree a
routedTree :: forall a. Ord a => RoutedSchedule a -> RoutedTree a
routedTree RoutedSchedule {ingress :: forall a. RoutedSchedule a -> a
ingress = a
member, routedSteps :: forall a. RoutedSchedule a -> Schedule a
routedSteps = Schedule a
steps} =
  a -> Schedule a -> RoutedTree a
forall a. Ord a => a -> Schedule a -> RoutedTree a
scheduleTree a
member Schedule a
steps

renderDecompositionTree :: [String] -> DecompositionTree -> String
renderDecompositionTree :: [String] -> DecompositionTree -> String
renderDecompositionTree [String]
members (TileTree Tree TileNode
tree) =
  [String] -> String
unlines ([Split] -> Tree TileNode -> [String]
renderTree [] Tree TileNode
tree)
  where
    tileIds :: [([Split], Int)]
tileIds = [[Split]] -> [Int] -> [([Split], Int)]
forall a b. [a] -> [b] -> [(a, b)]
zip (([Split] -> (Int, [Split])) -> [[Split]] -> [[Split]]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn [Split] -> (Int, [Split])
forall {t :: * -> *} {a}. Foldable t => t a -> (Int, t a)
pathKey ([Split] -> Tree TileNode -> [[Split]]
decompositionPaths [] Tree TileNode
tree)) [Int
0 ..]

    renderTree :: [Split] -> Tree TileNode -> [String]
renderTree [Split]
path (Tree TileNode
node [Tree TileNode]
kids) =
      [String] -> Int -> TileNode -> String
renderNumberedNode [String]
members ([([Split], Int)] -> [Split] -> Int
tileId [([Split], Int)]
tileIds [Split]
path) TileNode
node
        String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [Split] -> [Tree TileNode] -> [String]
renderChildren [Split]
path [Tree TileNode]
kids

    renderChildren :: [Split] -> [Tree TileNode] -> [String]
renderChildren [Split]
_ [] = []
    renderChildren [Split]
path [Tree TileNode]
kids =
      [[String]] -> [String]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
        [ String -> Tree TileNode -> [String]
renderBranch String
prefix Tree TileNode
child
        | (String
prefix, Tree TileNode
child) <- [Tree TileNode] -> [(String, Tree TileNode)]
forall a. [a] -> [(String, a)]
branchPrefixes [Tree TileNode]
kids
        ]
      where
        renderBranch :: String -> Tree TileNode -> [String]
renderBranch String
prefix Tree TileNode
child =
          case [Split] -> Tree TileNode -> [String]
renderTree ([Split]
path [Split] -> [Split] -> [Split]
forall a. [a] -> [a] -> [a]
++ TileNode -> [Split]
relationPath (Tree TileNode -> TileNode
forall a. Tree a -> a
treeLabel Tree TileNode
child)) Tree TileNode
child of
            [] -> []
            String
first : [String]
rest ->
              (String
prefix String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
first)
                String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [ShowS
continuation String
prefix String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
line | String
line <- [String]
rest]

    pathKey :: t a -> (Int, t a)
pathKey t a
path = (t a -> Int
forall a. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t a
path, t a
path)

renderHopTree :: [String] -> HopTree -> String
renderHopTree :: [String] -> HopTree -> String
renderHopTree [String]
members (TileTree Tree TileNode
tree) =
  (TileNode -> String) -> Tree TileNode -> String
forall a. (a -> String) -> Tree a -> String
renderTreeWith ([String] -> Tile -> String
renderTile [String]
members (Tile -> String) -> (TileNode -> Tile) -> TileNode -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TileNode -> Tile
tile) Tree TileNode
tree

renderSendTree :: [String] -> SendTree -> String
renderSendTree :: [String] -> SendTree -> String
renderSendTree [String]
members (SendTree Tree Tile
tree) =
  (Tile -> String) -> Tree Tile -> String
forall a. (a -> String) -> Tree a -> String
renderTreeWith ([String] -> Tile -> String
renderMember [String]
members) Tree Tile
tree

renderRoutedTree :: RoutedTree String -> String
renderRoutedTree :: RoutedTree String -> String
renderRoutedTree (RoutedTree Tree String
tree) =
  ShowS -> Tree String -> String
forall a. (a -> String) -> Tree a -> String
renderTreeWith ShowS
forall a. a -> a
id Tree String
tree

mapTree :: (a -> b) -> Tree a -> Tree b
mapTree :: forall a b. (a -> b) -> Tree a -> Tree b
mapTree a -> b
f (Tree a
label [Tree a]
kids) =
  Tree
    { treeLabel :: b
treeLabel = a -> b
f a
label,
      subtrees :: [Tree b]
subtrees = (Tree a -> Tree b) -> [Tree a] -> [Tree b]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> Tree a -> Tree b
forall a b. (a -> b) -> Tree a -> Tree b
mapTree a -> b
f) [Tree a]
kids
    }

unfoldTree :: (a -> [a]) -> a -> Tree a
unfoldTree :: forall a. (a -> [a]) -> a -> Tree a
unfoldTree a -> [a]
childrenOf a
label =
  Tree
    { treeLabel :: a
treeLabel = a
label,
      subtrees :: [Tree a]
subtrees = (a -> Tree a) -> [a] -> [Tree a]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> [a]) -> a -> Tree a
forall a. (a -> [a]) -> a -> Tree a
unfoldTree a -> [a]
childrenOf) (a -> [a]
childrenOf a
label)
    }

renderTreeWith :: (a -> String) -> Tree a -> String
renderTreeWith :: forall a. (a -> String) -> Tree a -> String
renderTreeWith a -> String
renderLabel Tree a
tree =
  [String] -> String
unlines (Tree a -> [String]
renderLines Tree a
tree)
  where
    renderLines :: Tree a -> [String]
renderLines (Tree a
label [Tree a]
kids) =
      a -> String
renderLabel a
label String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [Tree a] -> [String]
renderChildren [Tree a]
kids

    renderChildren :: [Tree a] -> [String]
renderChildren [] = []
    renderChildren [Tree a]
kids =
      [[String]] -> [String]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
        [ String -> Tree a -> [String]
renderBranch String
prefix Tree a
child
        | (String
prefix, Tree a
child) <- [Tree a] -> [(String, Tree a)]
forall a. [a] -> [(String, a)]
branchPrefixes [Tree a]
kids
        ]

    renderBranch :: String -> Tree a -> [String]
renderBranch String
prefix Tree a
child =
      case Tree a -> [String]
renderLines Tree a
child of
        [] -> []
        String
first : [String]
rest ->
          (String
prefix String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
first)
            String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [ShowS
continuation String
prefix String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
line | String
line <- [String]
rest]

branchPrefixes :: [a] -> [(String, a)]
branchPrefixes :: forall a. [a] -> [(String, a)]
branchPrefixes [] = []
branchPrefixes [a
x] = [(String
"└─ ", a
x)]
branchPrefixes (a
x : [a]
xs) = (String
"├─ ", a
x) (String, a) -> [(String, a)] -> [(String, a)]
forall a. a -> [a] -> [a]
: [a] -> [(String, a)]
forall a. [a] -> [(String, a)]
branchPrefixes [a]
xs

continuation :: String -> String
continuation :: ShowS
continuation String
"├─ " = String
"│  "
continuation String
"└─ " = String
"   "
continuation String
prefix = Int -> Char -> String
forall a. Int -> a -> [a]
replicate (String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
prefix) Char
' '

decompositionPaths :: [Split] -> Tree TileNode -> [[Split]]
decompositionPaths :: [Split] -> Tree TileNode -> [[Split]]
decompositionPaths [Split]
path (Tree TileNode
_ [Tree TileNode]
kids) =
  [Split]
path [Split] -> [[Split]] -> [[Split]]
forall a. a -> [a] -> [a]
: [[[Split]]] -> [[Split]]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Split] -> Tree TileNode -> [[Split]]
decompositionPaths ([Split]
path [Split] -> [Split] -> [Split]
forall a. [a] -> [a] -> [a]
++ TileNode -> [Split]
relationPath (Tree TileNode -> TileNode
forall a. Tree a -> a
treeLabel Tree TileNode
child)) Tree TileNode
child | Tree TileNode
child <- [Tree TileNode]
kids]

relationPath :: TileNode -> [Split]
relationPath :: TileNode -> [Split]
relationPath TileNode
node =
  case TileNode -> Relation
relation TileNode
node of
    Relation
Root -> []
    Anchor Split
split -> [Split
split]
    Sibling Split
split -> [Split
split]

tileId :: [([Split], Int)] -> [Split] -> Int
tileId :: [([Split], Int)] -> [Split] -> Int
tileId [([Split], Int)]
ids [Split]
path =
  case [Split] -> [([Split], Int)] -> Maybe Int
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup [Split]
path [([Split], Int)]
ids of
    Just Int
ident -> Int
ident
    Maybe Int
Nothing -> String -> Int
forall a. HasCallStack => String -> a
error String
"tileId: missing path"

renderNumberedNode :: [String] -> Int -> TileNode -> String
renderNumberedNode :: [String] -> Int -> TileNode -> String
renderNumberedNode [String]
members Int
ident TileNode
node =
  String
"T" String -> ShowS
forall a. [a] -> [a] -> [a]
++ Int -> String
forall a. Show a => a -> String
show Int
ident String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
" " String -> ShowS
forall a. [a] -> [a] -> [a]
++ [String] -> Tile -> String
renderTile [String]
members (TileNode -> Tile
tile TileNode
node)

renderTile :: [String] -> Tile -> String
renderTile :: [String] -> Tile -> String
renderTile [String]
members Tile
tile =
  String
"[" String -> ShowS
forall a. [a] -> [a] -> [a]
++ [String] -> String
unwords [[String]
members [String] -> Int -> String
forall a. HasCallStack => [a] -> Int -> a
!! Int
rank | Int
rank <- Tile -> [Int]
tileRanks Tile
tile] String -> ShowS
forall a. [a] -> [a] -> [a]
++ String
"]"

renderMember :: [String] -> Tile -> String
renderMember :: [String] -> Tile -> String
renderMember [String]
members Tile
tile =
  [String]
members [String] -> Int -> String
forall a. HasCallStack => [a] -> Int -> a
!! Tile -> Int
root Tile
tile