content: module Crux.Loops where
  
  import qualified Data.Map as Map
  
  data Tree a = Node { node :: a, children :: Forest a }
                      deriving Show
  
  type Forest a = [Tree a]
  
  findLoops :: Eq a => [a] -> Forest a
  findLoops = go []
    where
    go prev steps =
      case steps of
        [] -> reverse prev
        x : xs ->
          case break ((x ==) . node) prev of
            (inner,_:prev') -> go (Node x [] : Node x (reverse inner) : prev') xs
            _               -> go (Node x [] : prev) xs
  
  
  annotate :: Ord a => Forest a -> [ ([Int],a) ]
  annotate = go [] Map.empty [] []
    where
    go ns iters prevIters rest todo =
      case todo of
        [] -> rest
        tree : trees ->
          let a = node tree
              n = Map.findWithDefault 1 a iters
              ns' = n : ns
              rest' = go ns (Map.insert a (n+1) iters) prevIters rest trees
          in (reverse ns', a) :
                go ns' Map.empty (iters : prevIters) rest' (children tree)
  
  
  annotateLoops :: Ord a => [a] -> [ ([Int],a) ]
  annotateLoops = annotate . findLoops
  
  
  