-- This file is part of the 'term-rewriting' library. It is licensed
-- under an MIT license. See the accompanying 'LICENSE' file for details.
--
-- Authors: Martin Avanzini, Bertram Felgenhauer

module Data.Rewriting.Term.Ops (
    -- * Operations on Terms
    funs,
    funsDL,
    vars,
    varsDL,
    root,
    withArity,
    subtermAt,
    properSubterms,
    subterms,
    replaceAt,
    rename,
    -- * Predicates on Terms
    isVar,
    isFun,
    isGround,
    isLinear,
    isInstanceOf,
    isVariantOf,
) where

import Data.Rewriting.Pos
import Data.Rewriting.Term.Type as Term
import Data.Rewriting.Substitution.Match
import Data.Maybe
import qualified Data.MultiSet as MS

import Control.Monad (guard)

-- | Annotate each occurrence of a function symbol with its actual arity,
-- i.e., its number of arguments.
--
-- >>> withArity (Fun 'f' [Var 1, Fun 'f' []])
-- Fun ('f',2) [Var 1,Fun ('f',0) []]
withArity :: Term f v -> Term (f, Int) v
withArity :: forall f v. Term f v -> Term (f, Int) v
withArity = (v -> Term (f, Int) v)
-> (f -> [Term (f, Int) v] -> Term (f, Int) v)
-> Term f v
-> Term (f, Int) v
forall v a f. (v -> a) -> (f -> [a] -> a) -> Term f v -> a
Term.fold v -> Term (f, Int) v
forall f v. v -> Term f v
Var (\f
f [Term (f, Int) v]
ts -> (f, Int) -> [Term (f, Int) v] -> Term (f, Int) v
forall f v. f -> [Term f v] -> Term f v
Fun (f
f, [Term (f, Int) v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Term (f, Int) v]
ts) [Term (f, Int) v]
ts)

-- | Return the subterm at a given position.
subtermAt :: Term f v -> Pos -> Maybe (Term f v)
subtermAt :: forall f v. Term f v -> Pos -> Maybe (Term f v)
subtermAt Term f v
t [] = Term f v -> Maybe (Term f v)
forall a. a -> Maybe a
Just Term f v
t
subtermAt (Fun f
_ [Term f v]
ts) (Int
p:Pos
ps) | Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Term f v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Term f v]
ts = Term f v -> Pos -> Maybe (Term f v)
forall f v. Term f v -> Pos -> Maybe (Term f v)
subtermAt ([Term f v]
ts [Term f v] -> Int -> Term f v
forall a. HasCallStack => [a] -> Int -> a
!! Int
p) Pos
ps
subtermAt Term f v
_ Pos
_ = Maybe (Term f v)
forall a. Maybe a
Nothing

-- | Return the list of all proper subterms.
--
-- >>> properSubterms (Fun 'g' [Fun 'f' [Var 1], Fun 'f' [Var 1]])
-- [Fun 'f' [Var 1],Var 1,Fun 'f' [Var 1],Var 1]
properSubterms :: Term f v -> [Term f v]
properSubterms :: forall f v. Term f v -> [Term f v]
properSubterms (Var v
_) = []
properSubterms (Fun f
_ [Term f v]
ts) = (Term f v -> [Term f v]) -> [Term f v] -> [Term f v]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Term f v -> [Term f v]
forall f v. Term f v -> [Term f v]
subterms [Term f v]
ts

-- | Return the list of all subterms.
--
-- prop> subterms t = t : properSubterms t
subterms :: Term f v -> [Term f v]
subterms :: forall f v. Term f v -> [Term f v]
subterms Term f v
t = Term f v
t Term f v -> [Term f v] -> [Term f v]
forall a. a -> [a] -> [a]
: Term f v -> [Term f v]
forall f v. Term f v -> [Term f v]
properSubterms Term f v
t

-- NOTE: replaceAt and Context.ofTerm have the same recusion structure; is
-- there a nice higher-order function to abstract from it?

-- | replace a subterm at a given position.
replaceAt :: Term f v -> Pos -> Term f v -> Maybe (Term f v)
replaceAt :: forall f v. Term f v -> Pos -> Term f v -> Maybe (Term f v)
replaceAt Term f v
_ [] Term f v
t' = Term f v -> Maybe (Term f v)
forall a. a -> Maybe a
Just Term f v
t'
replaceAt (Fun f
f [Term f v]
ts) (Int
i:Pos
p) Term f v
t' = do
    Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [Term f v] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Term f v]
ts)
    let ([Term f v]
ts1, Term f v
t:[Term f v]
ts2) = Int -> [Term f v] -> ([Term f v], [Term f v])
forall a. Int -> [a] -> ([a], [a])
splitAt Int
i [Term f v]
ts
    Term f v
t <- Term f v -> Pos -> Term f v -> Maybe (Term f v)
forall f v. Term f v -> Pos -> Term f v -> Maybe (Term f v)
replaceAt Term f v
t Pos
p Term f v
t'
    Term f v -> Maybe (Term f v)
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Term f v -> Maybe (Term f v)) -> Term f v -> Maybe (Term f v)
forall a b. (a -> b) -> a -> b
$ f -> [Term f v] -> Term f v
forall f v. f -> [Term f v] -> Term f v
Fun f
f ([Term f v]
ts1 [Term f v] -> [Term f v] -> [Term f v]
forall a. [a] -> [a] -> [a]
++ Term f v
t Term f v -> [Term f v] -> [Term f v]
forall a. a -> [a] -> [a]
: [Term f v]
ts2)
replaceAt Term f v
_ Pos
_ Term f v
_ = Maybe (Term f v)
forall a. Maybe a
Nothing

-- | Return the list of all variables in the term, from left to right.
--
-- >>> vars (Fun 'g' [Var 3, Fun 'f' [Var 1, Var 2, Var 3]])
-- [3,1,2,3]
vars :: Term f v -> [v]
vars :: forall f v. Term f v -> [v]
vars = (Term f v -> [v] -> [v]) -> [v] -> Term f v -> [v]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Term f v -> [v] -> [v]
forall f v. Term f v -> [v] -> [v]
varsDL []

-- | Difference List version of 'vars'.
-- We have @varsDL t vs = vars t ++ vs@.

varsDL :: Term f v -> [v] -> [v]
varsDL :: forall f v. Term f v -> [v] -> [v]
varsDL = (v -> [v] -> [v])
-> (f -> [[v] -> [v]] -> [v] -> [v]) -> Term f v -> [v] -> [v]
forall v a f. (v -> a) -> (f -> [a] -> a) -> Term f v -> a
Term.fold (:) (([[v] -> [v]] -> [v] -> [v]) -> f -> [[v] -> [v]] -> [v] -> [v]
forall a b. a -> b -> a
const (([[v] -> [v]] -> [v] -> [v]) -> f -> [[v] -> [v]] -> [v] -> [v])
-> ([[v] -> [v]] -> [v] -> [v]) -> f -> [[v] -> [v]] -> [v] -> [v]
forall a b. (a -> b) -> a -> b
$ (([v] -> [v]) -> ([v] -> [v]) -> [v] -> [v])
-> ([v] -> [v]) -> [[v] -> [v]] -> [v] -> [v]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ([v] -> [v]) -> ([v] -> [v]) -> [v] -> [v]
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) [v] -> [v]
forall a. a -> a
id)


-- | Return the root symbol of the given term.
--
-- >>> root (Fun 'f' [Var 1, Fun 'g' []])
-- Right 'f'
--
-- >>> root (Var 1)
-- Left 1
root :: Term f v -> Either v f
root :: forall f v. Term f v -> Either v f
root (Fun f
f [Term f v]
_) = f -> Either v f
forall a b. b -> Either a b
Right f
f
root (Var v
v) = v -> Either v f
forall a b. a -> Either a b
Left v
v

-- | Return the list of all function symbols in the term, from left to right.
--
-- >>> funs (Fun 'f' [Var 3, Fun 'g' [Fun 'f' []]])
-- "fgf"
funs :: Term f v -> [f]
funs :: forall f v. Term f v -> [f]
funs = (Term f v -> [f] -> [f]) -> [f] -> Term f v -> [f]
forall a b c. (a -> b -> c) -> b -> a -> c
flip Term f v -> [f] -> [f]
forall f v. Term f v -> [f] -> [f]
funsDL []

-- | Difference List version of 'funs'.
-- We have @funsDL t vs = funs t ++ vs@.
funsDL :: Term f v -> [f] -> [f]
funsDL :: forall f v. Term f v -> [f] -> [f]
funsDL = (v -> [f] -> [f])
-> (f -> [[f] -> [f]] -> [f] -> [f]) -> Term f v -> [f] -> [f]
forall v a f. (v -> a) -> (f -> [a] -> a) -> Term f v -> a
Term.fold (([f] -> [f]) -> v -> [f] -> [f]
forall a b. a -> b -> a
const [f] -> [f]
forall a. a -> a
id) (\f
f [[f] -> [f]]
xs -> (f
ff -> [f] -> [f]
forall a. a -> [a] -> [a]
:) ([f] -> [f]) -> ([f] -> [f]) -> [f] -> [f]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (([f] -> [f]) -> ([f] -> [f]) -> [f] -> [f])
-> ([f] -> [f]) -> [[f] -> [f]] -> [f] -> [f]
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ([f] -> [f]) -> ([f] -> [f]) -> [f] -> [f]
forall b c a. (b -> c) -> (a -> b) -> a -> c
(.) [f] -> [f]
forall a. a -> a
id [[f] -> [f]]
xs)

-- | Return 'True' if the term is a variable, 'False' otherwise.
isVar :: Term f v -> Bool
isVar :: forall f v. Term f v -> Bool
isVar Var{} = Bool
True
isVar Fun{} = Bool
False

-- | Return 'True' if the term is a function application, 'False' otherwise.
isFun :: Term f v -> Bool
isFun :: forall f v. Term f v -> Bool
isFun Var{} = Bool
False
isFun Fun{} = Bool
True

-- | Check whether the term is a ground term, i.e., contains no variables.
isGround :: Term f v -> Bool
isGround :: forall f v. Term f v -> Bool
isGround = [v] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([v] -> Bool) -> (Term f v -> [v]) -> Term f v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term f v -> [v]
forall f v. Term f v -> [v]
vars

-- | Check whether the term is linear, i.e., contains each variable at most
-- once.
isLinear :: Ord v => Term f v -> Bool
isLinear :: forall v f. Ord v => Term f v -> Bool
isLinear = ((v, Int) -> Bool) -> [(v, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(v
_, Int
c) -> Int
c Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1) ([(v, Int)] -> Bool)
-> (Term f v -> [(v, Int)]) -> Term f v -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. MultiSet v -> [(v, Int)]
forall a. MultiSet a -> [(a, Int)]
MS.toOccurList (MultiSet v -> [(v, Int)])
-> (Term f v -> MultiSet v) -> Term f v -> [(v, Int)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [v] -> MultiSet v
forall a. Ord a => [a] -> MultiSet a
MS.fromList ([v] -> MultiSet v) -> (Term f v -> [v]) -> Term f v -> MultiSet v
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Term f v -> [v]
forall f v. Term f v -> [v]
vars

-- | Check whether the first term is an instance of the second term.
isInstanceOf :: (Eq f, Ord v, Ord v') => Term f v -> Term f v' -> Bool
isInstanceOf :: forall f v v'.
(Eq f, Ord v, Ord v') =>
Term f v -> Term f v' -> Bool
isInstanceOf Term f v
t Term f v'
u = Maybe (GSubst v' f v) -> Bool
forall a. Maybe a -> Bool
isJust (Term f v' -> Term f v -> Maybe (GSubst v' f v)
forall f v v'.
(Eq f, Ord v, Eq v') =>
Term f v -> Term f v' -> Maybe (GSubst v f v')
match Term f v'
u Term f v
t)

-- | Check whether two terms are variants of each other.
isVariantOf :: (Eq f, Ord v, Ord v') => Term f v -> Term f v' -> Bool
isVariantOf :: forall f v v'.
(Eq f, Ord v, Ord v') =>
Term f v -> Term f v' -> Bool
isVariantOf Term f v
t Term f v'
u = Term f v -> Term f v' -> Bool
forall f v v'.
(Eq f, Ord v, Ord v') =>
Term f v -> Term f v' -> Bool
isInstanceOf Term f v
t Term f v'
u Bool -> Bool -> Bool
&& Term f v' -> Term f v -> Bool
forall f v v'.
(Eq f, Ord v, Ord v') =>
Term f v -> Term f v' -> Bool
isInstanceOf Term f v'
u Term f v
t

-- | Rename the variables in a term.
--
-- >>> rename (+ 1) (Fun 'f' [Var 1, Fun 'g' [Var 2]])
-- (Fun 'f' [Var 2, Fun 'g' [Var 3]])
rename :: (v -> v') -> Term f v -> Term f v'
rename :: forall v v' f. (v -> v') -> Term f v -> Term f v'
rename = (f -> f) -> (v -> v') -> Term f v -> Term f v'
forall f f' v v'. (f -> f') -> (v -> v') -> Term f v -> Term f' v'
Term.map f -> f
forall a. a -> a
id