{-# LANGUAGE TemplateHaskell #-}

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
$c== :: Path -> Path -> Bool
== :: Path -> Path -> Bool
$c/= :: Path -> Path -> Bool
/= :: 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
$cshowsPrec :: Int -> Path -> ShowS
showsPrec :: Int -> Path -> ShowS
$cshow :: Path -> String
show :: Path -> String
$cshowList :: [Path] -> ShowS
showList :: [Path] -> ShowS
Show)

instance Ord Path where
  compare :: Path -> Path -> Ordering
compare (Path [Int]
p1) (Path [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 String
"Cannot take parent of empty path"
parentPath (Path (Int
_:[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 []     [Int]
ys      = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
ys
revPathDistance [Int]
xs     []      = [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
xs
revPathDistance (Int
x:[Int]
xs) (Int
y:[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 a. [a] -> 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 a. [a] -> 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 Path
p1 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 :: forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
HFix f i -> Int -> Maybe (E (HFix f))
getChild (Term f (HFix f) i
t) 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)) :=> First (E (HFix f))
forall m (a :: * -> *). Monoid m => (a :=> m) -> f a :=> m
forall (h :: (* -> *) -> * -> *) m (a :: * -> *).
(HFoldable h, Monoid m) =>
(a :=> m) -> h a :=> m
hfoldMap Numbered (HFix f) i -> First (E (HFix f))
Numbered (HFix f) :=> First (E (HFix f))
forall (g :: (* -> *) -> * -> *) j.
Numbered (HFix g) j -> First (E (HFix g))
eqP (f (HFix f) i -> f (Numbered (HFix f)) i
f (HFix f) :-> f (Numbered (HFix f))
forall (f :: (* -> *) -> * -> *) (a :: * -> *).
HTraversable f =>
f a :-> f (Numbered a)
number f (HFix f) i
t))
  where
    eqP :: forall g j. Numbered (HFix g) j -> First (E (HFix g))
    eqP :: forall (g :: (* -> *) -> * -> *) j.
Numbered (HFix g) j -> First (E (HFix g))
eqP (Numbered Int
j HFix g j
x)
          | Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
j    = Maybe (E (HFix g)) -> First (E (HFix g))
forall a. Maybe a -> First a
First (E (HFix g) -> Maybe (E (HFix g))
forall a. a -> Maybe a
Just (HFix g j -> E (HFix g)
forall (f :: * -> *) i. f i -> E f
E HFix g j
x))
          | Bool
otherwise = Maybe (E (HFix g)) -> First (E (HFix g))
forall a. Maybe a -> First a
First Maybe (E (HFix g))
forall a. Maybe a
Nothing


rewriteChild :: forall m f i. (HTraversable f, Applicative m) => Int -> HFix f i -> (forall j. HFix f j -> m (HFix f j)) -> m (HFix f i)
rewriteChild :: forall (m :: * -> *) (f :: (* -> *) -> * -> *) 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 (Term f (Cxt NoHole f (K ())) i
t) forall j. HFix f j -> m (HFix f j)
f = f (Cxt NoHole f (K ())) i -> Cxt NoHole f (K ()) i
forall (f :: (* -> *) -> * -> *) h (a :: * -> *) i.
f (Cxt h f a) i -> Cxt h f a i
Term (f (Cxt NoHole f (K ())) i -> Cxt NoHole f (K ()) i)
-> m (f (Cxt NoHole f (K ())) i) -> m (Cxt NoHole f (K ()) i)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> NatM m (Numbered (Cxt NoHole f (K ()))) (Cxt NoHole f (K ()))
-> NatM
     m (f (Numbered (Cxt NoHole f (K ())))) (f (Cxt NoHole f (K ())))
forall (f :: * -> *) (a :: * -> *) (b :: * -> *).
Applicative f =>
NatM f a b -> NatM f (f a) (f b)
forall (t :: (* -> *) -> * -> *) (f :: * -> *) (a :: * -> *)
       (b :: * -> *).
(HTraversable t, Applicative f) =>
NatM f a b -> NatM f (t a) (t b)
htraverse Numbered (Cxt NoHole f (K ())) i -> m (HFix f i)
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
f (Cxt NoHole f (K ())) :-> f (Numbered (Cxt NoHole f (K ())))
forall (f :: (* -> *) -> * -> *) (a :: * -> *).
HTraversable f =>
f a :-> f (Numbered a)
number f (Cxt NoHole f (K ())) i
t)
  where
    rw :: forall j. Numbered (HFix f) j -> m (HFix f j)
    rw :: NatM m (Numbered (Cxt NoHole f (K ()))) (Cxt NoHole f (K ()))
rw (Numbered Int
j 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 a. a -> m a
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 :: forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
[Int] -> HFix f i -> Maybe (E (HFix f))
followRevPath []     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 (Int
i:[Int]
is) 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] -> Cxt NoHole f (K ()) 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 :: forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
Path -> HFix f i -> [E (HFix f)]
getAncestors (Path [Int]
p) 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 :: forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
[Int] -> HFix f i -> [E (HFix f)]
go []     HFix f i
t = [HFix f i -> E (HFix f)
forall (f :: * -> *) i. f i -> E f
E HFix f i
t]
    go (Int
i:[Int]
is) 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 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
                    Maybe (E (HFix f))
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 :: forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
Path -> HFix f i -> Maybe (E (HFix f))
followPath Path
p 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 (f :: (* -> *) -> * -> *) (m :: * -> *) i.
(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)
f HFix f i
t (Path [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 HFix f j -> m (HFix f j)
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 (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 []     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 forall j. HFix f j -> m (HFix f j)
f (Int
i:[Int]
is) HFix f i
x = Int
-> HFix f i -> (forall j. HFix f j -> m (HFix f j)) -> m (HFix f i)
forall (m :: * -> *) (f :: (* -> *) -> * -> *) 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 HFix f j -> m (HFix f j)
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 :: forall (f :: (* -> *) -> * -> *) f'.
HTraversable f =>
Alg (f :&: Label) (K ([Int] -> Map Label Path))
pathAlg (:&:) 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
$ \[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 a (f :: (* -> *) -> * -> *) (e :: * -> *) l.
Annotated a f =>
f e l -> a
forall (e :: * -> *) l. (:&:) f Label e l -> Label
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
(:&:) f Label (K ([Int] -> Map Label Path))
:-> (:&:) f Label (Numbered (K ([Int] -> Map Label Path)))
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 [Int]
path = (Numbered (K ([Int] -> Map Label Path)) :=> Map Label Path)
-> (:&:) f Label (Numbered (K ([Int] -> Map Label Path)))
   :=> Map Label Path
forall m (a :: * -> *).
Monoid m =>
(a :=> m) -> (:&:) f Label a :=> m
forall (h :: (* -> *) -> * -> *) m (a :: * -> *).
(HFoldable h, Monoid m) =>
(a :=> m) -> h a :=> m
hfoldMap (\(Numbered Int
i K ([Int] -> Map Label Path) 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 :: forall (f :: (* -> *) -> * -> *) i.
HTraversable f =>
HFixLab f i -> Map Label Path
getPaths 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))
-> HFix (f :&: Label) :-> K ([Int] -> Map Label Path)
forall (f :: (* -> *) -> * -> *) (a :: * -> *).
HFunctor f =>
Alg f a -> HFix f :-> a
cata (:&:) f Label (K ([Int] -> Map Label Path)) i
-> K ([Int] -> Map Label Path) i
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 (f :: (* -> *) -> * -> *) l.
HTraversable f =>
(forall i. HFix f i -> Bool)
-> HFix f l -> Path -> Maybe (E (HFix f))
searchParent forall i. HFix f i -> Bool
f HFix f l
prog 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 HFix f i -> Bool
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)