module GF.Command.TreeOperations (
  treeOp,
  allTreeOps,
  treeChunks
  ) where

import PGF(Expr,PGF,CId,compute,mkApp,unApp,unapply,unMeta,exprSize,exprFunctions)
import Data.List

type TreeOp = [Expr] -> [Expr]

treeOp :: PGF -> String -> Maybe (Either TreeOp (CId -> TreeOp))
treeOp :: PGF -> String -> Maybe (Either TreeOp (CId -> TreeOp))
treeOp PGF
pgf String
f = ((String, Either TreeOp (CId -> TreeOp))
 -> Either TreeOp (CId -> TreeOp))
-> Maybe (String, Either TreeOp (CId -> TreeOp))
-> Maybe (Either TreeOp (CId -> TreeOp))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (String, Either TreeOp (CId -> TreeOp))
-> Either TreeOp (CId -> TreeOp)
forall a b. (a, b) -> b
snd (Maybe (String, Either TreeOp (CId -> TreeOp))
 -> Maybe (Either TreeOp (CId -> TreeOp)))
-> Maybe (String, Either TreeOp (CId -> TreeOp))
-> Maybe (Either TreeOp (CId -> TreeOp))
forall a b. (a -> b) -> a -> b
$ String
-> [(String, (String, Either TreeOp (CId -> TreeOp)))]
-> Maybe (String, Either TreeOp (CId -> TreeOp))
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup String
f ([(String, (String, Either TreeOp (CId -> TreeOp)))]
 -> Maybe (String, Either TreeOp (CId -> TreeOp)))
-> [(String, (String, Either TreeOp (CId -> TreeOp)))]
-> Maybe (String, Either TreeOp (CId -> TreeOp))
forall a b. (a -> b) -> a -> b
$ PGF -> [(String, (String, Either TreeOp (CId -> TreeOp)))]
allTreeOps PGF
pgf

allTreeOps :: PGF -> [(String,(String,Either TreeOp (CId -> TreeOp)))]
allTreeOps :: PGF -> [(String, (String, Either TreeOp (CId -> TreeOp)))]
allTreeOps PGF
pgf = [
   (String
"compute",(String
"compute by using semantic definitions (def)",
      TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. a -> Either a b
Left  (TreeOp -> Either TreeOp (CId -> TreeOp))
-> TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. (a -> b) -> a -> b
$ (Expr -> Expr) -> TreeOp
forall a b. (a -> b) -> [a] -> [b]
map (PGF -> Expr -> Expr
compute PGF
pgf))),
   (String
"largest",(String
"sort trees from largest to smallest, in number of nodes",
      TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. a -> Either a b
Left  (TreeOp -> Either TreeOp (CId -> TreeOp))
-> TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. (a -> b) -> a -> b
$ TreeOp
largest)),
   (String
"nub",(String
"remove duplicate trees",
      TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. a -> Either a b
Left  (TreeOp -> Either TreeOp (CId -> TreeOp))
-> TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. (a -> b) -> a -> b
$ TreeOp
forall a. Eq a => [a] -> [a]
nub)),
   (String
"smallest",(String
"sort trees from smallest to largest, in number of nodes",
      TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. a -> Either a b
Left  (TreeOp -> Either TreeOp (CId -> TreeOp))
-> TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. (a -> b) -> a -> b
$ TreeOp
smallest)),
   (String
"subtrees",(String
"return all fully applied subtrees (stopping at abstractions), by default sorted from the largest",
      TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. a -> Either a b
Left  (TreeOp -> Either TreeOp (CId -> TreeOp))
-> TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. (a -> b) -> a -> b
$ (Expr -> [Expr]) -> TreeOp
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr -> [Expr]
subtrees)),
   (String
"funs",(String
"return all fun functions appearing in the tree, with duplications",
      TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. a -> Either a b
Left  (TreeOp -> Either TreeOp (CId -> TreeOp))
-> TreeOp -> Either TreeOp (CId -> TreeOp)
forall a b. (a -> b) -> a -> b
$ \[Expr]
es -> [CId -> [Expr] -> Expr
mkApp CId
f [] | Expr
e <- [Expr]
es, CId
f <- Expr -> [CId]
exprFunctions Expr
e]))
  ]

largest :: [Expr] -> [Expr]
largest :: TreeOp
largest = TreeOp
forall a. [a] -> [a]
reverse TreeOp -> TreeOp -> TreeOp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TreeOp
smallest

smallest :: [Expr] -> [Expr]
smallest :: TreeOp
smallest = (Expr -> Expr -> Ordering) -> TreeOp
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (\Expr
t Expr
u -> Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
compare (Expr -> Int
exprSize Expr
t) (Expr -> Int
exprSize Expr
u))

treeChunks :: Expr -> [Expr]
treeChunks :: Expr -> [Expr]
treeChunks = (Bool, [Expr]) -> [Expr]
forall a b. (a, b) -> b
snd ((Bool, [Expr]) -> [Expr])
-> (Expr -> (Bool, [Expr])) -> Expr -> [Expr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> (Bool, [Expr])
cks where
  cks :: Expr -> (Bool, [Expr])
cks Expr
t = 
    case Expr -> (Expr, [Expr])
unapply Expr
t of
      (Expr
t, [Expr]
ts) -> case Expr -> Maybe Int
unMeta Expr
t of
                   Just Int
_  -> (Bool
False,(Expr -> [Expr]) -> TreeOp
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ((Bool, [Expr]) -> [Expr]
forall a b. (a, b) -> b
snd ((Bool, [Expr]) -> [Expr])
-> (Expr -> (Bool, [Expr])) -> Expr -> [Expr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> (Bool, [Expr])
cks) [Expr]
ts)
                   Maybe Int
Nothing -> case [(Bool, [Expr])] -> ([Bool], [[Expr]])
forall a b. [(a, b)] -> ([a], [b])
unzip ((Expr -> (Bool, [Expr])) -> [Expr] -> [(Bool, [Expr])]
forall a b. (a -> b) -> [a] -> [b]
map Expr -> (Bool, [Expr])
cks [Expr]
ts) of
                                ([Bool]
bs,[[Expr]]
_) | [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
and [Bool]
bs      -> (Bool
True, [Expr
t])
                                ([Bool]
_,[[Expr]]
cts)              -> (Bool
False,[[Expr]] -> [Expr]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [[Expr]]
cts)

subtrees :: Expr -> [Expr]
subtrees :: Expr -> [Expr]
subtrees Expr
t = Expr
t Expr -> TreeOp
forall a. a -> [a] -> [a]
: case Expr -> Maybe (CId, [Expr])
unApp Expr
t of
  Just (CId
f,[Expr]
ts) -> (Expr -> [Expr]) -> TreeOp
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr -> [Expr]
subtrees [Expr]
ts
  Maybe (CId, [Expr])
_ -> []  -- don't go under abstractions