module Crux.Loops where

import qualified Data.Map as Map

data Tree a = Node { forall a. Tree a -> a
node :: a, forall a. Tree a -> Forest a
children :: 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
$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

type Forest a = [Tree a]

findLoops :: Eq a => [a] -> Forest a
findLoops :: forall a. Eq a => [a] -> Forest a
findLoops = [Tree a] -> [a] -> [Tree a]
forall {a}. Eq a => [Tree a] -> [a] -> [Tree a]
go []
  where
  go :: [Tree a] -> [a] -> [Tree a]
go [Tree a]
prev [a]
steps =
    case [a]
steps of
      [] -> [Tree a] -> [Tree a]
forall a. [a] -> [a]
reverse [Tree a]
prev
      a
x : [a]
xs ->
        case (Tree a -> Bool) -> [Tree a] -> ([Tree a], [Tree a])
forall a. (a -> Bool) -> [a] -> ([a], [a])
break ((a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
==) (a -> Bool) -> (Tree a -> a) -> Tree a -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree a -> a
forall a. Tree a -> a
node) [Tree a]
prev of
          ([Tree a]
inner,Tree a
_:[Tree a]
prev') -> [Tree a] -> [a] -> [Tree a]
go (a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
x [] Tree a -> [Tree a] -> [Tree a]
forall a. a -> [a] -> [a]
: a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
x ([Tree a] -> [Tree a]
forall a. [a] -> [a]
reverse [Tree a]
inner) Tree a -> [Tree a] -> [Tree a]
forall a. a -> [a] -> [a]
: [Tree a]
prev') [a]
xs
          ([Tree a], [Tree a])
_               -> [Tree a] -> [a] -> [Tree a]
go (a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
x [] Tree a -> [Tree a] -> [Tree a]
forall a. a -> [a] -> [a]
: [Tree a]
prev) [a]
xs


annotate :: Ord a => Forest a -> [ ([Int],a) ]
annotate :: forall a. Ord a => Forest a -> [([Int], a)]
annotate = [Int]
-> Map a Int
-> [Map a Int]
-> [([Int], a)]
-> [Tree a]
-> [([Int], a)]
forall {b} {a}.
(Ord b, Num a) =>
[a] -> Map b a -> [Map b a] -> [([a], b)] -> [Tree b] -> [([a], b)]
go [] Map a Int
forall k a. Map k a
Map.empty [] []
  where
  go :: [a] -> Map b a -> [Map b a] -> [([a], b)] -> [Tree b] -> [([a], b)]
go [a]
ns Map b a
iters [Map b a]
prevIters [([a], b)]
rest [Tree b]
todo =
    case [Tree b]
todo of
      [] -> [([a], b)]
rest
      Tree b
tree : [Tree b]
trees ->
        let a :: b
a = Tree b -> b
forall a. Tree a -> a
node Tree b
tree
            n :: a
n = a -> b -> Map b a -> a
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault a
1 b
a Map b a
iters
            ns' :: [a]
ns' = a
n a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
ns
            rest' :: [([a], b)]
rest' = [a] -> Map b a -> [Map b a] -> [([a], b)] -> [Tree b] -> [([a], b)]
go [a]
ns (b -> a -> Map b a -> Map b a
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert b
a (a
na -> a -> a
forall a. Num a => a -> a -> a
+a
1) Map b a
iters) [Map b a]
prevIters [([a], b)]
rest [Tree b]
trees
        in ([a] -> [a]
forall a. [a] -> [a]
reverse [a]
ns', b
a) ([a], b) -> [([a], b)] -> [([a], b)]
forall a. a -> [a] -> [a]
:
              [a] -> Map b a -> [Map b a] -> [([a], b)] -> [Tree b] -> [([a], b)]
go [a]
ns' Map b a
forall k a. Map k a
Map.empty (Map b a
iters Map b a -> [Map b a] -> [Map b a]
forall a. a -> [a] -> [a]
: [Map b a]
prevIters) [([a], b)]
rest' (Tree b -> [Tree b]
forall a. Tree a -> Forest a
children Tree b
tree)


annotateLoops :: Ord a => [a] -> [ ([Int],a) ]
annotateLoops :: forall a. Ord a => [a] -> [([Int], a)]
annotateLoops = Forest a -> [([Int], a)]
forall a. Ord a => Forest a -> [([Int], a)]
annotate (Forest a -> [([Int], a)])
-> ([a] -> Forest a) -> [a] -> [([Int], a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> Forest a
forall a. Eq a => [a] -> Forest a
findLoops