TNumber.hs

{-# LANGUAGE TupleSections #-}
{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}
module TNumber where

import Text.Pretty.Simple
import Control.Monad.State

-- https://stackoverflow.com/questions/44784899/how-to-write-function-for-n-ary-tree-traversal-in-haskell?rq=1

data NT a = N a [NT a] deriving (Show, Functor, Foldable, Traversable)


{-
aNumber   :: a                -- thing to number
          -> Int              -- number to start from
          -> ( (a, Int)       -- numbered thing
             , Int            -- next available number afterwards
             )
aNumber a i = ((a, i), i + 1)


ntNumber  :: NT a             -- thing to number
          -> Int              -- number to start from
          -> ( NT (a, Int)    -- numbered thing
             , Int            -- next available number afterwards
             )
ntNumber (N a ants) i0  = (N ai aints, i2) where
  (ai,    i1) = aNumber   a    i0
  (aints, i2) = ntsNumber ants i1

ntsNumber :: [NT a]           -- thing to number
          -> Int              -- number to start from
          -> ( [NT (a, Int)]  -- numbered thing
             , Int            -- next available number afterwards
             )
ntsNumber []           i  = ([], i)
ntsNumber (ant : ants) i0 = (aint : aints, i2) where
  (aint,  i1) = ntNumber  ant  i0
  (aints, i2) = ntsNumber ants i1
-}


ntree :: NT String
ntree = N "eric" [N "lea" [N "kristy" [],N "pedro" [] ,N "rafael" []],N "anna" [],N "bety" []]

myTree :: NT [String]
myTree = N ["a", "b", "c"] [N ["d", "e"] [], N ["f"] []]

-- Second Attempt: Numbering and Threading
--

type Numbering output = Int -> (output, Int)

aNumber :: t -> Numbering (t, Int)
aNumber a = steady (a, ) $$ next

ntNumber  :: NT a -> Numbering (NT (a, Int))
ntNumber (N a ants) = steady N $$ aNumber a $$ listNumber ntNumber ants

ntNumber' :: (a -> Numbering b) -> NT a -> Numbering (NT b)
ntNumber' na (N a ants) = steady N $$ na a $$ listNumber (ntNumber' na) ants

ntsNumber :: [NT a] -> Numbering [NT (a, Int)]
-- ntsNumber []           = steady []
-- ntsNumber (ant : ants) = steady (:) $$ ntNumber ant $$ ntsNumber ants
ntsNumber = listNumber ntNumber


next :: Numbering Int
next i = (i, i + 1)

next' :: State Int Int
next' = get <* modify (1+)


steady :: thing -> Numbering thing
steady x i = (x, i)

($$) :: Numbering (a -> b) -> Numbering a -> Numbering b
(fn $$ an) i0 = (f a, i2) where
  (f, i1) = fn i0
  (a, i2) = an i1
infixl 2 $$


listNumber :: (a -> Numbering b) -> [a] -> Numbering [b]
listNumber na = foldr (\ a -> ($$) (steady (:) $$ na a)) (steady [])

result :: NT a -> NT (a, Int)
result nt = evalState (traverse (\ a ->  gets (a,) <* modify (1+)) nt) 0