Flattened to Tree structure

flattened-to-tree.hs

-- こういうフラットな配列から木構造を作るのってどうやったらいいんだろう？
-- https://twitter.com/suin/status/1390626013814476800

{-# LANGUAGE DeriveFunctor, UndecidableInstances, LambdaCase #-}

module Main where

import Data.Functor.Classes
import Data.List

----------------------------------------

-- Q. How to change flattened structure into Tree structure?

inputNodes :: [NodeF a]
inputNodes = [
    NodeF [1] "title 1" [],
    NodeF [1, 1] "title 1-1" [],
    NodeF [1, 1, 1] "title 1-1-1" [],
    NodeF [1, 1, 2] "title 1-1-2" [],
    NodeF [1, 2] "title 1-2" [],
    NodeF [2] "title 2" []
  ]

outputNodes :: [Node]
outputNodes = [
    mkNode [1] "title 1" [
      mkNode [1] "title 1-1" [
        mkNode [1] "title 1-1-1" [],
        mkNode [2] "title 1-1-2" []
      ],
      mkNode [2] "title 1-2" []
    ],
    mkNode [2] "title 2" []
  ]

----------------------------------------

newtype Fix f = In { out :: f (Fix f) }

instance Show (f (Fix f)) => Show (Fix f) where
  showsPrec p (In f) = showsPrec p f

instance Eq1 f => Eq (Fix f) where
  In a == In b = eq1 a b

cata :: Functor f => (f a -> a) -> Fix f -> a
cata g = g . fmap (cata g) . out

ana :: Functor f => (a -> f a) -> a -> Fix f
ana g = In . fmap (ana g) . g

----------------------------------------

data NodeF a = NodeF [Int] String [a] deriving (Functor, Show, Eq)

instance Eq1 NodeF where
  liftEq _ (NodeF _ _ []) (NodeF _ _ []) = True
  liftEq _ (NodeF _ _ []) (NodeF _ _ (_:_)) = False
  liftEq _ (NodeF _ _ (_:_)) (NodeF _ _ []) = False
  liftEq eq (NodeF _ _ (x:xs)) (NodeF _ _ (y:ys)) = eq x y && liftEq eq xs ys

type Node = Fix NodeF

mkNode :: [Int] -> String -> [Node] -> Node
mkNode keys title children = In (NodeF keys title children)

getKeys :: Node -> [Int]
getKeys (In (NodeF ks _ _)) = ks

----------------------------------------

-- Step 1: Make a nested Tree structure using Anamorphism.

-- Example:
-- coalg (NodeF [1,2,3] "title 1-2-3" [])
--   = NodeF [1] "" [ NodeF [2, 3] "title 1-2-3" [] ]
coalg :: NodeF a -> NodeF (NodeF a)
coalg (NodeF (k : []) title _) = NodeF [k] title []
coalg (NodeF (k : ks) title _) = NodeF [k] "" [(NodeF ks title [])]

-- nestedInputNode :: Node
-- nestedInputNode = ana coalg $ NodeF [1, 1, 2] "title 1-1-2" []
-- -- == NodeF [1] "" [ NodeF [1] "" [ NodeF [2] "title 1-1-2" [] ] ]

nestedInputNodes :: [Node]
nestedInputNodes = fmap (ana coalg) inputNodes

----------------------------------------

-- Step 2: Fold `nestedInputNodes` using Catamorphism.

data ListF a b = Nil | Cons a b deriving (Functor, Show, Eq)

type NodeList = Fix (ListF Node)

toList' :: [a] -> Fix (ListF a)
toList' [] = In Nil
toList' (a : as) = In (Cons a (toList' as))

fromList' :: Fix (ListF a) -> [a]
fromList' (In Nil) = []
fromList' (In (Cons a list)) = a : (fromList' list)

reducer :: Node -> [Node] -> [Node]
reducer n@(In (NodeF ks title children)) ns =
  case elemIndex ks (fmap getKeys ns) of
    Just i -> fmap (f i) $ zip [0..] ns
      where
        f i (j, n@(In (NodeF ks2 title2 children2))) =
          if i == j
            then In (NodeF ks (title ++ title2) (maybe children2 (\x -> reducer x children2) (safeHead children)))
            else n

        safeHead [] = Nothing
        safeHead (x:xs) = Just x

    Nothing -> [n] ++ ns

alg :: ListF Node NodeList -> NodeList
alg = \case
  Nil -> In Nil
  Cons n list -> toList' $ reducer n (fromList' list)

answer :: [Node]
answer = fromList' (cata alg (toList' nestedInputNodes))

----------------------------------------

-- Step 2b: Fold `inputNodes` using Catamorphism with running Anamorphism at the same time.

alg' :: ListF (NodeF a) NodeList -> NodeList
alg' = \case
  Nil -> In Nil
  Cons n list -> toList' $ reducer (ana coalg n) (fromList' list)

answer' :: [Node]
answer' = fromList' (cata alg' (toList' inputNodes))

----------------------------------------

main :: IO ()
main = do
  print $ foldr reducer [] nestedInputNodes -- using foldr
  print $ answer -- using catamorphism
  print $ answer' -- using catamorphism

  print $ answer == outputNodes -- True
  print $ answer' == outputNodes -- True