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