{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeOperators #-}

module Cubix.Language.Parametric.Path
  (
    Path
  , emptyPath
  , isEmptyPath
  , parentPath
  , pathDistance
  , getChild
  , rewriteChild
  , followPath
  , rewriteAtPathM
  , getAncestors
  , getPaths
  , searchParent
  ) where


import Data.Comp.Multi ( Cxt(..), Alg, cata, K(..), E(..), runE, (:&:), hfoldMap, HTraversable(..), HFix )
import Data.Comp.Multi.Mapping ( Numbered(..), number )

import Data.List ( find )
import Data.Map ( Map )
import qualified Data.Map as Map
import Data.Monoid ( First(..) )

import Cubix.Language.Info

import Cubix.Sin.Compdata.Annotation ( getAnn' )

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

newtype Path = Path {Path -> [Int]
getPath :: [Int]}
 deriving (Path -> Path -> Bool
(Path -> Path -> Bool) -> (Path -> Path -> Bool) -> Eq Path
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Path -> Path -> Bool
$c/= :: Path -> Path -> Bool
== :: Path -> Path -> Bool
$c== :: Path -> Path -> Bool
Eq, Int -> Path -> ShowS
[Path] -> ShowS
Path -> String
(Int -> Path -> ShowS)
-> (Path -> String) -> ([Path] -> ShowS) -> Show Path
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Path] -> ShowS
$cshowList :: [Path] -> ShowS
show :: Path -> String
$cshow :: Path -> String
showsPrec :: Int -> Path -> ShowS
$cshowsPrec :: Int -> Path -> ShowS
Show)

instance Ord Path where
  compare :: Path -> Path -> Ordering
compare (Path p1 :: [Int]
p1) (Path p2 :: [Int]
p2) = [Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
p1 [Int] -> [Int] -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` [Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
p2

emptyPath :: Path
emptyPath :: Path
emptyPath = [Int] -> Path
Path []

isEmptyPath :: Path -> Bool
isEmptyPath :: Path -> Bool
isEmptyPath = ([Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
==[]) ([Int] -> Bool) -> (Path -> [Int]) -> Path -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Path -> [Int]
getPath

parentPath :: Path -> Path
parentPath :: Path -> Path
parentPath (Path [])    = String -> Path
forall a. HasCallStack => String -> a
error "Cannot take parent of empty path"
parentPath (Path (_:p :: [Int]
p)) = [Int] -> Path
Path [Int]
p

-- A crappy but easy-to-implement metric for the distance between two nodes,
-- as represented by their paths
--
-- Note that lists of statements are represented as cons cells, so this will be
-- approximately linear in the line distance for things drawn from the same block
revPathDistance :: [Int] -> [Int] -> Int
revPathDistance :: [Int] -> [Int] -> Int
revPathDistance []     ys :: [Int]
ys      = [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ys
revPathDistance xs :: [Int]
xs     []      = [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
xs
revPathDistance (x :: Int
x:xs :: [Int]
xs) (y :: Int
y:ys :: [Int]
ys)
                   | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y    = [Int] -> [Int] -> Int
revPathDistance [Int]
xs [Int]
ys
                   | Bool
otherwise = [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Int
xInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
xs) Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [Int] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Int
yInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
ys)

pathDistance :: Path -> Path -> Int
pathDistance :: Path -> Path -> Int
pathDistance p1 :: Path
p1 p2 :: Path
p2 = [Int] -> [Int] -> Int
revPathDistance ([Int] -> [Int]
forall a. [a] -> [a]
reverse ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Path -> [Int]
getPath Path
p1) ([Int] -> [Int]
forall a. [a] -> [a]
reverse ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Path -> [Int]
getPath Path
p2)

getChild :: (HTraversable f) => HFix f i -> Int -> Maybe (E (HFix f))
getChild :: HFix f i -> Int -> Maybe (E (HFix f))
getChild (Term t :: f (HFix f) i
t) i :: Int
i = First (E (HFix f)) -> Maybe (E (HFix f))
forall a. First a -> Maybe a
getFirst ((Numbered (HFix f) :=> First (E (HFix f)))
-> f (Numbered (HFix f)) i -> First (E (HFix f))
forall (h :: (* -> *) -> * -> *) m (a :: * -> *).
(HFoldable h, Monoid m) =>
(a :=> m) -> h a :=> m
hfoldMap Numbered (HFix f) :=> First (E (HFix f))
forall (f :: * -> *) i. Numbered f i -> First (E f)
eqP (f (HFix f) i -> f (Numbered (HFix f)) i
forall (f :: (* -> *) -> * -> *) (a :: * -> *).
HTraversable f =>
f a :-> f (Numbered a)
number f (HFix f) i
t))
  where
    eqP :: Numbered f i -> First (E f)
eqP (Numbered j :: Int
j x :: f i
x)
          | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j    = Maybe (E f) -> First (E f)
forall a. Maybe a -> First a
First (E f -> Maybe (E f)
forall a. a -> Maybe a
Just (f i -> E f
forall (f :: * -> *) i. f i -> E f
E f i
x))
          | Bool
otherwise = Maybe (E f) -> First (E f)
forall a. Maybe a -> First a
First Maybe (E f)
forall a. Maybe a
Nothing


rewriteChild :: (HTraversable f, Applicative m) => Int -> HFix f i -> (forall j. HFix f j -> m (HFix f j)) -> m (HFix f i)
rewriteChild :: Int
-> HFix f i -> (forall j. HFix f j -> m (HFix f j)) -> m (HFix f i)
rewriteChild i :: Int
i (Term t :: f (Cxt NoHole f (K ())) i
t) f :: forall j. HFix f j -> m (HFix f j)
f = f (Cxt NoHole f (K ())) i -> HFix f i
forall (f :: (* -> *) -> * -> *) h (a :: * -> *) i.
f (Cxt h f a) i -> Cxt h f a i
Term (f (Cxt NoHole f (K ())) i -> HFix f i)
-> m (f (Cxt NoHole f (K ())) i) -> m (HFix f i)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NatM m (Numbered (Cxt NoHole f (K ()))) (Cxt NoHole f (K ()))
-> f (Numbered (Cxt NoHole f (K ()))) i
-> m (f (Cxt NoHole f (K ())) i)
forall (t :: (* -> *) -> * -> *) (f :: * -> *) (a :: * -> *)
       (b :: * -> *).
(HTraversable t, Applicative f) =>
NatM f a b -> NatM f (t a) (t b)
htraverse NatM m (Numbered (Cxt NoHole f (K ()))) (Cxt NoHole f (K ()))
rw (f (Cxt NoHole f (K ())) i -> f (Numbered (Cxt NoHole f (K ()))) i
forall (f :: (* -> *) -> * -> *) (a :: * -> *).
HTraversable f =>
f a :-> f (Numbered a)
number f (Cxt NoHole f (K ())) i
t)
  where
    rw :: Numbered (Cxt NoHole f (K ())) j -> m (HFix f j)
rw (Numbered j :: Int
j x :: HFix f j
x)
          | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j    = HFix f j -> m (HFix f j)
forall j. HFix f j -> m (HFix f j)
f HFix f j
x
          | Bool
otherwise = HFix f j -> m (HFix f j)
forall (f :: * -> *) a. Applicative f => a -> f a
pure HFix f j
x

followRevPath :: (HTraversable f) => [Int] -> HFix f i -> Maybe (E (HFix f))
followRevPath :: [Int] -> HFix f i -> Maybe (E (HFix f))
followRevPath []     t :: HFix f i
t = E (HFix f) -> Maybe (E (HFix f))
forall a. a -> Maybe a
Just (HFix f i -> E (HFix f)
forall (f :: * -> *) i. f i -> E f
E HFix f i
t)
followRevPath (i :: Int
i:is :: [Int]
is) t :: HFix f i
t = (HFix f :=> Maybe (E (HFix f))) -> E (HFix f) -> Maybe (E (HFix f))
forall (f :: * -> *) b. (f :=> b) -> E f -> b
runE ([Int] -> HFix f i -> Maybe (E (HFix f))
forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
[Int] -> HFix f i -> Maybe (E (HFix f))
followRevPath [Int]
is) (E (HFix f) -> Maybe (E (HFix f)))
-> Maybe (E (HFix f)) -> Maybe (E (HFix f))
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HFix f i -> Int -> Maybe (E (HFix f))
forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
HFix f i -> Int -> Maybe (E (HFix f))
getChild HFix f i
t Int
i

getAncestors :: (HTraversable f) => Path -> HFix f i -> [E (HFix f)]
getAncestors :: Path -> HFix f i -> [E (HFix f)]
getAncestors (Path p :: [Int]
p) t :: HFix f i
t = [Int] -> HFix f i -> [E (HFix f)]
forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
[Int] -> HFix f i -> [E (HFix f)]
go ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
p) HFix f i
t
  where
    go :: (HTraversable f) => [Int] -> HFix f i -> [E (HFix f)]
    go :: [Int] -> HFix f i -> [E (HFix f)]
go []     t :: HFix f i
t = [HFix f i -> E (HFix f)
forall (f :: * -> *) i. f i -> E f
E HFix f i
t]
    go (i :: Int
i:is :: [Int]
is) t :: HFix f i
t = case HFix f i -> Int -> Maybe (E (HFix f))
forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
HFix f i -> Int -> Maybe (E (HFix f))
getChild HFix f i
t Int
i of
                    Just (E x :: HFix f i
x) -> (HFix f i -> E (HFix f)
forall (f :: * -> *) i. f i -> E f
E HFix f i
t) E (HFix f) -> [E (HFix f)] -> [E (HFix f)]
forall a. a -> [a] -> [a]
: [Int] -> HFix f i -> [E (HFix f)]
forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
[Int] -> HFix f i -> [E (HFix f)]
go [Int]
is HFix f i
x
                    Nothing    -> [HFix f i -> E (HFix f)
forall (f :: * -> *) i. f i -> E f
E HFix f i
t]

followPath :: (HTraversable f) => Path -> HFix f i -> Maybe (E (HFix f))
followPath :: Path -> HFix f i -> Maybe (E (HFix f))
followPath p :: Path
p t :: HFix f i
t = [Int] -> HFix f i -> Maybe (E (HFix f))
forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
[Int] -> HFix f i -> Maybe (E (HFix f))
followRevPath ([Int] -> [Int]
forall a. [a] -> [a]
reverse ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Path -> [Int]
getPath Path
p) HFix f i
t

rewriteAtPathM :: (HTraversable f, Applicative m) => (forall j. HFix f j -> m (HFix f j)) -> HFix f i -> Path -> m (HFix f i)
rewriteAtPathM :: (forall j. HFix f j -> m (HFix f j))
-> HFix f i -> Path -> m (HFix f i)
rewriteAtPathM f :: forall j. HFix f j -> m (HFix f j)
f t :: HFix f i
t (Path p :: [Int]
p) = (forall j. HFix f j -> m (HFix f j))
-> [Int] -> HFix f i -> m (HFix f i)
forall (f :: (* -> *) -> * -> *) (m :: * -> *) i.
(HTraversable f, Applicative m) =>
(forall j. HFix f j -> m (HFix f j))
-> [Int] -> HFix f i -> m (HFix f i)
go forall j. HFix f j -> m (HFix f j)
f ([Int] -> [Int]
forall a. [a] -> [a]
reverse [Int]
p) HFix f i
t
  where
    go :: (HTraversable f, Applicative m) =>  (forall j. HFix f j -> m (HFix f j)) -> [Int] -> HFix f i -> m (HFix f i)
    go :: (forall j. HFix f j -> m (HFix f j))
-> [Int] -> HFix f i -> m (HFix f i)
go f :: forall j. HFix f j -> m (HFix f j)
f []     x :: HFix f i
x = HFix f i -> m (HFix f i)
forall j. HFix f j -> m (HFix f j)
f HFix f i
x
    go f :: forall j. HFix f j -> m (HFix f j)
f (i :: Int
i:is :: [Int]
is) x :: HFix f i
x = Int
-> HFix f i -> (forall j. HFix f j -> m (HFix f j)) -> m (HFix f i)
forall (f :: (* -> *) -> * -> *) (m :: * -> *) i.
(HTraversable f, Applicative m) =>
Int
-> HFix f i -> (forall j. HFix f j -> m (HFix f j)) -> m (HFix f i)
rewriteChild Int
i HFix f i
x ((forall j. HFix f j -> m (HFix f j))
-> [Int] -> HFix f j -> m (HFix f j)
forall (f :: (* -> *) -> * -> *) (m :: * -> *) i.
(HTraversable f, Applicative m) =>
(forall j. HFix f j -> m (HFix f j))
-> [Int] -> HFix f i -> m (HFix f i)
go forall j. HFix f j -> m (HFix f j)
f [Int]
is)

pathAlg :: forall f f'. (HTraversable f) => Alg (f :&: Label) (K ([Int] -> Map Label Path))
pathAlg :: Alg (f :&: Label) (K ([Int] -> Map Label Path))
pathAlg t :: (:&:) f Label (K ([Int] -> Map Label Path)) i
t = ([Int] -> Map Label Path) -> K ([Int] -> Map Label Path) i
forall a i. a -> K a i
K (([Int] -> Map Label Path) -> K ([Int] -> Map Label Path) i)
-> ([Int] -> Map Label Path) -> K ([Int] -> Map Label Path) i
forall a b. (a -> b) -> a -> b
$ \path :: [Int]
path -> Label -> Path -> Map Label Path -> Map Label Path
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Label
lab ([Int] -> Path
Path [Int]
path) ([Int] -> Map Label Path
childPaths [Int]
path)
  where
    lab :: Label
lab = (:&:) f Label (K ([Int] -> Map Label Path)) i -> Label
forall (f :: (* -> *) -> * -> *) a (e :: * -> *) l.
Annotated f a =>
f e l -> a
getAnn' (:&:) f Label (K ([Int] -> Map Label Path)) i
t
    tInd :: (:&:) f Label (Numbered (K ([Int] -> Map Label Path))) i
tInd = (:&:) f Label (K ([Int] -> Map Label Path)) i
-> (:&:) f Label (Numbered (K ([Int] -> Map Label Path))) i
forall (f :: (* -> *) -> * -> *) (a :: * -> *).
HTraversable f =>
f a :-> f (Numbered a)
number (:&:) f Label (K ([Int] -> Map Label Path)) i
t
    childPaths :: [Int] -> Map Label Path
childPaths path :: [Int]
path = (Numbered (K ([Int] -> Map Label Path)) :=> Map Label Path)
-> (:&:) f Label (Numbered (K ([Int] -> Map Label Path))) i
-> Map Label Path
forall (h :: (* -> *) -> * -> *) m (a :: * -> *).
(HFoldable h, Monoid m) =>
(a :=> m) -> h a :=> m
hfoldMap (\(Numbered i kf) -> (K ([Int] -> Map Label Path) i -> [Int] -> Map Label Path
forall a i. K a i -> a
unK K ([Int] -> Map Label Path) i
kf) (Int
iInt -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:[Int]
path)) (:&:) f Label (Numbered (K ([Int] -> Map Label Path))) i
tInd

getPaths :: (HTraversable f) => HFixLab f i -> Map Label Path
getPaths :: HFixLab f i -> Map Label Path
getPaths t :: HFixLab f i
t = K ([Int] -> Map Label Path) i -> [Int] -> Map Label Path
forall a i. K a i -> a
unK (Alg (f :&: Label) (K ([Int] -> Map Label Path))
-> HFixLab f i -> K ([Int] -> Map Label Path) i
forall (f :: (* -> *) -> * -> *) (a :: * -> *).
HFunctor f =>
Alg f a -> HFix f :-> a
cata Alg (f :&: Label) (K ([Int] -> Map Label Path))
forall (f :: (* -> *) -> * -> *) f'.
HTraversable f =>
Alg (f :&: Label) (K ([Int] -> Map Label Path))
pathAlg HFixLab f i
t) []

searchParent :: (HTraversable f) => (forall i. HFix f i -> Bool) -> HFix f l -> Path -> Maybe (E (HFix f))
searchParent :: (forall i. HFix f i -> Bool)
-> HFix f l -> Path -> Maybe (E (HFix f))
searchParent f :: forall i. HFix f i -> Bool
f prog :: HFix f l
prog path :: Path
path = (E (HFix f) -> Bool) -> [E (HFix f)] -> Maybe (E (HFix f))
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Maybe a
find ((forall i. HFix f i -> Bool) -> E (HFix f) -> Bool
forall (f :: * -> *) b. (f :=> b) -> E f -> b
runE forall i. HFix f i -> Bool
f) ([E (HFix f)] -> [E (HFix f)]
forall a. [a] -> [a]
reverse ([E (HFix f)] -> [E (HFix f)]) -> [E (HFix f)] -> [E (HFix f)]
forall a b. (a -> b) -> a -> b
$ Path -> HFix f l -> [E (HFix f)]
forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
Path -> HFix f i -> [E (HFix f)]
getAncestors Path
path HFix f l
prog)