{-# LANGUAGE LambdaCase          #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Data.Record.TH.CodeGen.Tree (
    -- * Trees and forests
    Tree(..)
  , Forest(..)
    -- * Catamorphisms
  , Cata(..)
  , tree
  , forest
    -- * Dealing with @ghc@'s tuple size limit
  , TupleLimit(..)
  , nest
  , mkTupleE
  , mkTupleT
  , mkTupleP
  ) where

import Language.Haskell.TH

import Data.Record.Internal.TH.Util

{-------------------------------------------------------------------------------
  Trees and forests
-------------------------------------------------------------------------------}

-- | Trees with values at the leaves
data Tree   a = Leaf a | Branch (Forest 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
showList :: [Tree a] -> ShowS
$cshowList :: forall a. Show a => [Tree a] -> ShowS
show :: Tree a -> String
$cshow :: forall a. Show a => Tree a -> String
showsPrec :: Int -> Tree a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Tree a -> ShowS
Show)
data Forest a = Forest [Tree a]            deriving (Int -> Forest a -> ShowS
[Forest a] -> ShowS
Forest a -> String
(Int -> Forest a -> ShowS)
-> (Forest a -> String) -> ([Forest a] -> ShowS) -> Show (Forest a)
forall a. Show a => Int -> Forest a -> ShowS
forall a. Show a => [Forest a] -> ShowS
forall a. Show a => Forest a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Forest a] -> ShowS
$cshowList :: forall a. Show a => [Forest a] -> ShowS
show :: Forest a -> String
$cshow :: forall a. Show a => Forest a -> String
showsPrec :: Int -> Forest a -> ShowS
$cshowsPrec :: forall a. Show a => Int -> Forest a -> ShowS
Show)

{-------------------------------------------------------------------------------
  Catamorphisms

  Unlike regular folds, these catamorphisms are structure preserving.
  See "Dealing with Large Bananas", by Ralf Lämmel et al
-------------------------------------------------------------------------------}

data Cata a b = Cata {
      Cata a b -> a -> b
leaf   :: a -> b
    , Cata a b -> [b] -> b
branch :: [b] -> b
    }

tree :: Cata a b -> Tree a -> b
tree :: Cata a b -> Tree a -> b
tree Cata a b
cata (Leaf   a
a)  = Cata a b -> a -> b
forall a b. Cata a b -> a -> b
leaf   Cata a b
cata a
a
tree Cata a b
cata (Branch Forest a
as) = Cata a b -> Forest a -> b
forall a b. Cata a b -> Forest a -> b
forest Cata a b
cata Forest a
as

forest :: Cata a b -> Forest a -> b
forest :: Cata a b -> Forest a -> b
forest Cata a b
cata (Forest [Tree a]
ts) = Cata a b -> [b] -> b
forall a b. Cata a b -> [b] -> b
branch Cata a b
cata ((Tree a -> b) -> [Tree a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map (Cata a b -> Tree a -> b
forall a b. Cata a b -> Tree a -> b
tree Cata a b
cata) [Tree a]
ts)

{-------------------------------------------------------------------------------
  Nesting
-------------------------------------------------------------------------------}

-- | Observe @ghc@'s tuple length
--
-- Haskell has a limit of 62 fields per tuple. Here we take an arbitrary
-- list and turn it into a nested tuple that preserves this limit.
--
-- Example: if we reduce the limit to @2@, we get the following nestings,
-- for lengths @[1..10]@:
--
-- >     A
-- >    (A, A)
-- >   ((A, A),  A)
-- >   ((A, A), (A, A))
-- >  (((A, A), (A, A)),   A)
-- >  (((A, A), (A, A)),  (A, A))
-- >  (((A, A), (A, A)), ((A, A),  A))
-- >  (((A, A), (A, A)), ((A, A), (A, A)))
-- > ((((A, A), (A, A)), ((A, A), (A, A))),  A)
-- > ((((A, A), (A, A)), ((A, A), (A, A))), (A, A))
nest :: TupleLimit -> [a] -> Forest a
nest :: TupleLimit -> [a] -> Forest a
nest TupleLimit
mLimit = [Tree a] -> Forest a
forall a. [Tree a] -> Forest a
go ([Tree a] -> Forest a) -> ([a] -> [Tree a]) -> [a] -> Forest a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Tree a) -> [a] -> [Tree a]
forall a b. (a -> b) -> [a] -> [b]
map a -> Tree a
forall a. a -> Tree a
Leaf
  where
    go :: [Tree a] -> Forest a
    go :: [Tree a] -> Forest a
go [Tree a]
ts | [Tree a] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Tree a]
ts Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
limit = [Tree a] -> Forest a
forall a. [Tree a] -> Forest a
Forest [Tree a]
ts
          | Bool
otherwise         = [Tree a] -> Forest a
forall a. [Tree a] -> Forest a
go (([Tree a] -> Tree a) -> [[Tree a]] -> [Tree a]
forall a b. (a -> b) -> [a] -> [b]
map (Forest a -> Tree a
forall a. Forest a -> Tree a
Branch (Forest a -> Tree a)
-> ([Tree a] -> Forest a) -> [Tree a] -> Tree a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Tree a] -> Forest a
forall a. [Tree a] -> Forest a
Forest) (Int -> [Tree a] -> [[Tree a]]
forall a. Int -> [a] -> [[a]]
chunk Int
limit [Tree a]
ts))

    limit :: Int
    limit :: Int
limit = case TupleLimit
mLimit of
              TupleLimit
DefaultGhcTupleLimit -> Int
62
              MaxTupleElems Int
n      -> Int
n

-- | Maximum number of elements in a tuple
data TupleLimit =
    -- | Default maximum number of elements in a tuple in ghc (62)
    DefaultGhcTupleLimit

    -- | Explicit specified liit
  | MaxTupleElems Int

{-------------------------------------------------------------------------------
  Constructing nested types/values/patterns
-------------------------------------------------------------------------------}

-- | Construct tuple type
mkTupleT :: forall a. (a -> Q Type) -> Forest a -> Q Type
mkTupleT :: (a -> Q Type) -> Forest a -> Q Type
mkTupleT a -> Q Type
f = Cata a (Q Type) -> Forest a -> Q Type
forall a b. Cata a b -> Forest a -> b
forest Cata a (Q Type)
cata
  where
    cata :: Cata a (Q Type)
    cata :: Cata a (Q Type)
cata = Cata :: forall a b. (a -> b) -> ([b] -> b) -> Cata a b
Cata {
          leaf :: a -> Q Type
leaf   = a -> Q Type
f
        , branch :: [Q Type] -> Q Type
branch = \case [Q Type
t] -> Q Type
t
                         [Q Type]
ts  -> Q Type -> [Q Type] -> Q Type
appsT (Int -> Q Type
tupleT ([Q Type] -> Int
forall (t :: Type -> Type) a. Foldable t => t a -> Int
length [Q Type]
ts)) [Q Type]
ts
        }

-- | Construct tuple expression
mkTupleE :: forall a. (a -> Q Exp) -> Forest a -> Q Exp
mkTupleE :: (a -> Q Exp) -> Forest a -> Q Exp
mkTupleE a -> Q Exp
f = Cata a (Q Exp) -> Forest a -> Q Exp
forall a b. Cata a b -> Forest a -> b
forest Cata a (Q Exp)
cata
  where
     cata :: Cata a (Q Exp)
     cata :: Cata a (Q Exp)
cata = Cata :: forall a b. (a -> b) -> ([b] -> b) -> Cata a b
Cata {
           leaf :: a -> Q Exp
leaf   = a -> Q Exp
f
         , branch :: [Q Exp] -> Q Exp
branch = \case [Q Exp
e] -> Q Exp
e
                          [Q Exp]
es  -> [Q Exp] -> Q Exp
tupE [Q Exp]
es
         }

-- | Construct tuple pattern
mkTupleP :: forall a. (a -> Q Pat) -> Forest a -> Q Pat
mkTupleP :: (a -> Q Pat) -> Forest a -> Q Pat
mkTupleP a -> Q Pat
f = Cata a (Q Pat) -> Forest a -> Q Pat
forall a b. Cata a b -> Forest a -> b
forest Cata a (Q Pat)
cata
  where
    cata :: Cata a (Q Pat)
    cata :: Cata a (Q Pat)
cata = Cata :: forall a b. (a -> b) -> ([b] -> b) -> Cata a b
Cata {
          leaf :: a -> Q Pat
leaf   = a -> Q Pat
f
        , branch :: [Q Pat] -> Q Pat
branch = \case [Q Pat
p] -> Q Pat
p
                         [Q Pat]
ps  -> [Q Pat] -> Q Pat
tupP [Q Pat]
ps
        }

{-------------------------------------------------------------------------------
  Internal auxiliary
-------------------------------------------------------------------------------}

chunk :: Int -> [a] -> [[a]]
chunk :: Int -> [a] -> [[a]]
chunk Int
n = [a] -> [[a]]
forall a. [a] -> [[a]]
go
  where
    go :: [a] -> [[a]]
    go :: [a] -> [[a]]
go [] = []
    go [a]
xs = let ([a]
firstChunk, [a]
rest) = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
n [a]
xs in [a]
firstChunk [a] -> [[a]] -> [[a]]
forall a. a -> [a] -> [a]
: [a] -> [[a]]
forall a. [a] -> [[a]]
go [a]
rest