-- |
-- Module      : Test.Speculate.Expr.Core
-- Copyright   : (c) 2016-2024 Rudy Matela
-- License     : 3-Clause BSD  (see the file LICENSE)
-- Maintainer  : Rudy Matela <rudy@matela.com.br>
--
-- This module is part of Speculate.
--
-- This module reexports 'Data.Express' along with some extra utilities.
module Test.Speculate.Expr.Core
  ( module Data.Express
  , module Data.Express.Utils.Typeable

  -- * Order
  , compareLexicographicallyBy
  , compareComplexityThenIndex

  -- * Properties
  , isConstantNamed
  , unrepeatedVars
  , isAssignment

  -- * Assigning
  , Binds

  -- * Matching
  , unify
  , unification
  , unificationC
  , isCanonInstanceOf
  , hasCanonInstanceOf

  -- * Commuting
  , commutations
  )
where

import Data.Express
import Data.Express.Utils.Typeable
import Test.Speculate.Utils
import Data.Monoid ((<>))
import Data.Functor ((<$>)) -- for GHC <= 7.8
import Data.Maybe (mapMaybe, listToMaybe)

-- | Lexicographical comparison of 'Expr's
--   where variables < constants < applications and
--   constants are disambiguated by the given function.
compareLexicographicallyBy :: (Expr -> Expr -> Ordering) -> Expr -> Expr -> Ordering
compareLexicographicallyBy :: (Expr -> Expr -> Ordering) -> Expr -> Expr -> Ordering
compareLexicographicallyBy Expr -> Expr -> Ordering
compareConstants  =  Expr -> Expr -> Ordering
cmp
  where
  (Expr
f :$ Expr
x) cmp :: Expr -> Expr -> Ordering
`cmp` (Expr
g :$ Expr
y)  =  Expr
f  Expr -> Expr -> Ordering
`cmp` Expr
g Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Expr
x Expr -> Expr -> Ordering
`cmp` Expr
y
  (Expr
_ :$ Expr
_) `cmp` Expr
_         =  Ordering
GT
  Expr
_        `cmp` (Expr
_ :$ Expr
_)  =  Ordering
LT
  Expr
e1 `cmp` Expr
e2  =  case (Expr -> Bool
isVar Expr
e1, Expr -> Bool
isVar Expr
e2) of
    (Bool
True,  Bool
True)  -> let Value String
n1 Dynamic
_ = Expr
e1
                          Value String
n2 Dynamic
_ = Expr
e2
                      in Expr -> TypeRep
typ Expr
e1 TypeRep -> TypeRep -> Ordering
`compareTy` Expr -> TypeRep
typ Expr
e2
                      Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
n1 Int -> Int -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
n2
                      Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> String
n1 String -> String -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` String
n2
    (Bool
False, Bool
True)  -> Ordering
GT
    (Bool
True,  Bool
False) -> Ordering
LT
    (Bool
False, Bool
False) -> Expr
e1 Expr -> Expr -> Ordering
`compareConstants` Expr
e2

-- | Comparison of 'Expr's:
--
-- 1. 'compareComplexity' from 'Data.Express'
-- 2. 'lexicompareBy' index of the given list
compareComplexityThenIndex :: [Expr] -> Expr -> Expr -> Ordering
compareComplexityThenIndex :: [Expr] -> Expr -> Expr -> Ordering
compareComplexityThenIndex [Expr]
as  =  Expr -> Expr -> Ordering
compareComplexity (Expr -> Expr -> Ordering)
-> (Expr -> Expr -> Ordering) -> Expr -> Expr -> Ordering
forall a. Semigroup a => a -> a -> a
<> (Expr -> Expr -> Ordering) -> Expr -> Expr -> Ordering
compareLexicographicallyBy Expr -> Expr -> Ordering
cmp
  where
  Expr
e1 cmp :: Expr -> Expr -> Ordering
`cmp` Expr
e2 | Expr -> Int
arity Expr
e1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 Bool -> Bool -> Bool
&& Expr -> Int
arity Expr
e2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 = Ordering
LT
  Expr
e1 `cmp` Expr
e2 | Expr -> Int
arity Expr
e1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
0 Bool -> Bool -> Bool
&& Expr -> Int
arity Expr
e2 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Ordering
GT
  Expr
e1 `cmp` Expr
e2 = [Expr] -> Expr -> Expr -> Ordering
forall a. Eq a => [a] -> a -> a -> Ordering
compareIndex [Expr]
as Expr
e1 Expr
e2 Ordering -> Ordering -> Ordering
forall a. Semigroup a => a -> a -> a
<> Expr
e1 Expr -> Expr -> Ordering
forall a. Ord a => a -> a -> Ordering
`compare` Expr
e2

countVars :: Expr -> [(Expr,Int)]
countVars :: Expr -> [(Expr, Int)]
countVars Expr
e = (Expr -> (Expr, Int)) -> [Expr] -> [(Expr, Int)]
forall a b. (a -> b) -> [a] -> [b]
map (\Expr
e' -> (Expr
e',[Expr] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Expr] -> Int) -> ([Expr] -> [Expr]) -> [Expr] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Expr -> Bool) -> [Expr] -> [Expr]
forall a. (a -> Bool) -> [a] -> [a]
filter (Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
e') ([Expr] -> Int) -> [Expr] -> Int
forall a b. (a -> b) -> a -> b
$ Expr -> [Expr]
vars Expr
e)) ([Expr] -> [(Expr, Int)]) -> [Expr] -> [(Expr, Int)]
forall a b. (a -> b) -> a -> b
$ Expr -> [Expr]
nubVars Expr
e

unrepeatedVars :: Expr -> Bool
unrepeatedVars :: Expr -> Bool
unrepeatedVars = ((Expr, Int) -> Bool) -> [(Expr, Int)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (\(Expr
_,Int
n) -> Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
1) ([(Expr, Int)] -> Bool) -> (Expr -> [(Expr, Int)]) -> Expr -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> [(Expr, Int)]
countVars

-- Is this espression an assignment of a variable to a value?
isAssignment :: Expr -> Bool
isAssignment :: Expr -> Bool
isAssignment ((Value String
"==" Dynamic
_ :$ Expr
e1) :$ Expr
e2) = Expr -> Bool
isVar Expr
e1 Bool -> Bool -> Bool
|| Expr -> Bool
isVar Expr
e2
isAssignment Expr
_ = Bool
False

isConstantNamed :: Expr -> String -> Bool
e :: Expr
e@(Value String
n' Dynamic
_) isConstantNamed :: Expr -> String -> Bool
`isConstantNamed` String
n = Expr -> Bool
isConst Expr
e Bool -> Bool -> Bool
&& String
n' String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
n
Expr
_ `isConstantNamed` String
_ = Bool
False

type Binds = [(Expr,Expr)]

unify :: Expr -> Expr -> Maybe Expr
unify :: Expr -> Expr -> Maybe Expr
unify Expr
e1 Expr
e2 = (Expr
e1 Expr -> [(Expr, Expr)] -> Expr
//-) ([(Expr, Expr)] -> Expr) -> Maybe [(Expr, Expr)] -> Maybe Expr
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Expr -> Expr -> Maybe [(Expr, Expr)]
unification Expr
e1 Expr
e2

unification :: Expr -> Expr -> Maybe [(Expr,Expr)]
unification :: Expr -> Expr -> Maybe [(Expr, Expr)]
unification = Expr -> Expr -> Maybe [(Expr, Expr)]
naiveUnification

findBind :: Expr -> Expr -> Either Bool (Expr,Expr)
findBind :: Expr -> Expr -> Either Bool (Expr, Expr)
findBind Expr
e1         Expr
e2          |  Expr -> TypeRep
typ Expr
e1 TypeRep -> TypeRep -> Bool
forall a. Eq a => a -> a -> Bool
/= Expr -> TypeRep
typ Expr
e2  =  Bool -> Either Bool (Expr, Expr)
forall a b. a -> Either a b
Left Bool
False
                                |  Expr
e1 Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
e2          =  Bool -> Either Bool (Expr, Expr)
forall a b. a -> Either a b
Left Bool
True
                                |  Expr -> Bool
isVar Expr
e1          =  (Expr, Expr) -> Either Bool (Expr, Expr)
forall a b. b -> Either a b
Right (Expr
e1,Expr
e2)
                                |  Expr -> Bool
isVar Expr
e2          =  (Expr, Expr) -> Either Bool (Expr, Expr)
forall a b. b -> Either a b
Right (Expr
e2,Expr
e1)
findBind (Expr
f1 :$ Expr
x1) (Expr
f2 :$ Expr
x2)  =  case Expr -> Expr -> Either Bool (Expr, Expr)
findBind Expr
f1 Expr
f2 of
                                   Left Bool
True -> Expr -> Expr -> Either Bool (Expr, Expr)
findBind Expr
x1 Expr
x2
                                   Either Bool (Expr, Expr)
r         -> Either Bool (Expr, Expr)
r
findBind Expr
e1         Expr
e2          =  Bool -> Either Bool (Expr, Expr)
forall a b. a -> Either a b
Left (Expr
e1 Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
e2)

-- NOTE: there are faster methods for unification.
naiveUnification :: Expr -> Expr -> Maybe [(Expr,Expr)]
naiveUnification :: Expr -> Expr -> Maybe [(Expr, Expr)]
naiveUnification Expr
e1' Expr
e2' = Expr -> Expr -> [(Expr, Expr)] -> Maybe [(Expr, Expr)]
uu Expr
e1' Expr
e2' []
  where
  uu :: Expr -> Expr -> Binds -> Maybe Binds
  uu :: Expr -> Expr -> [(Expr, Expr)] -> Maybe [(Expr, Expr)]
uu Expr
e1' Expr
e2' [(Expr, Expr)]
bs' =
    case Expr
-> Expr -> [(Expr, Expr)] -> Maybe (Expr, Expr, [(Expr, Expr)])
u Expr
e1' Expr
e2' [(Expr, Expr)]
bs' of
      Maybe (Expr, Expr, [(Expr, Expr)])
Nothing -> Maybe [(Expr, Expr)]
forall a. Maybe a
Nothing
      Just (Expr
e1,Expr
e2,[(Expr, Expr)]
bs) ->
        if Expr
e1' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
e1 Bool -> Bool -> Bool
&& Expr
e2' Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
e2
        then [(Expr, Expr)] -> Maybe [(Expr, Expr)]
forall a. a -> Maybe a
Just [(Expr, Expr)]
bs
        else Expr -> Expr -> [(Expr, Expr)] -> Maybe [(Expr, Expr)]
uu Expr
e1 Expr
e2 [(Expr, Expr)]
bs
  u :: Expr -> Expr -> Binds -> Maybe (Expr,Expr,Binds)
  u :: Expr
-> Expr -> [(Expr, Expr)] -> Maybe (Expr, Expr, [(Expr, Expr)])
u Expr
e1 Expr
e2 [(Expr, Expr)]
bs =
    case Expr -> Expr -> Either Bool (Expr, Expr)
findBind Expr
e1 Expr
e2 of
    Left Bool
False -> Maybe (Expr, Expr, [(Expr, Expr)])
forall a. Maybe a
Nothing
    Left Bool
True  -> (Expr, Expr, [(Expr, Expr)]) -> Maybe (Expr, Expr, [(Expr, Expr)])
forall a. a -> Maybe a
Just (Expr
e1,Expr
e2,[(Expr, Expr)]
bs)
    Right (Expr
ex,Expr
e) ->
      if Expr
ex Expr -> Expr -> Bool
`isSubexprOf` Expr
e
      then Maybe (Expr, Expr, [(Expr, Expr)])
forall a. Maybe a
Nothing
      else (Expr, Expr, [(Expr, Expr)]) -> Maybe (Expr, Expr, [(Expr, Expr)])
forall a. a -> Maybe a
Just ( Expr
e1 Expr -> [(Expr, Expr)] -> Expr
//- [(Expr
ex,Expr
e)]
                , Expr
e2 Expr -> [(Expr, Expr)] -> Expr
//- [(Expr
ex,Expr
e)]
                , (Expr
ex,Expr
e)(Expr, Expr) -> [(Expr, Expr)] -> [(Expr, Expr)]
forall a. a -> [a] -> [a]
:[(Expr
ex',Expr
e' Expr -> [(Expr, Expr)] -> Expr
//- [(Expr
ex,Expr
e)]) | (Expr
ex',Expr
e') <- [(Expr, Expr)]
bs]
                )

isCanonInstanceOf :: Expr -> Expr -> Bool
Expr
e1 isCanonInstanceOf :: Expr -> Expr -> Bool
`isCanonInstanceOf` Expr
e2 =
  case Expr
e1 Expr -> Expr -> Maybe [(Expr, Expr)]
`match` Expr
e2 of
    Maybe [(Expr, Expr)]
Nothing -> Bool
False
    Just [(Expr, Expr)]
xs -> ((Expr, Expr) -> Expr) -> [(Expr, Expr)] -> Bool
forall b a. Ord b => (a -> b) -> [a] -> Bool
strictlyOrderedOn (Expr, Expr) -> Expr
forall a b. (a, b) -> b
snd (((Expr, Expr) -> Expr) -> [(Expr, Expr)] -> [(Expr, Expr)]
forall b a. Ord b => (a -> b) -> [a] -> [a]
sortOn (Expr, Expr) -> Expr
forall a b. (a, b) -> a
fst [(Expr, Expr)]
xs)

hasCanonInstanceOf :: Expr -> Expr -> Bool
Expr
e1           hasCanonInstanceOf :: Expr -> Expr -> Bool
`hasCanonInstanceOf` Expr
e2 | Expr
e1   Expr -> Expr -> Bool
`isCanonInstanceOf` Expr
e2 = Bool
True
(Expr
e1f :$ Expr
e1x) `hasCanonInstanceOf` Expr
e2 | Expr
e1f Expr -> Expr -> Bool
`hasCanonInstanceOf` Expr
e2 Bool -> Bool -> Bool
||
                                       Expr
e1x Expr -> Expr -> Bool
`hasCanonInstanceOf` Expr
e2 = Bool
True
Expr
_            `hasCanonInstanceOf` Expr
_                                = Bool
False

unificationC :: [Expr] -> Expr -> Expr -> Maybe [(Expr,Expr)]
unificationC :: [Expr] -> Expr -> Expr -> Maybe [(Expr, Expr)]
unificationC [Expr]
cos Expr
e  =  [[(Expr, Expr)]] -> Maybe [(Expr, Expr)]
forall a. [a] -> Maybe a
listToMaybe
                    ([[(Expr, Expr)]] -> Maybe [(Expr, Expr)])
-> (Expr -> [[(Expr, Expr)]]) -> Expr -> Maybe [(Expr, Expr)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  (Expr -> Maybe [(Expr, Expr)]) -> [Expr] -> [[(Expr, Expr)]]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Expr
e Expr -> Expr -> Maybe [(Expr, Expr)]
`unification`)
                    ([Expr] -> [[(Expr, Expr)]])
-> (Expr -> [Expr]) -> Expr -> [[(Expr, Expr)]]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  [Expr] -> Expr -> [Expr]
commutations [Expr]
cos

commutations :: [Expr] -> Expr -> [Expr]
commutations :: [Expr] -> Expr -> [Expr]
commutations [Expr]
cos  =  Expr -> [Expr]
cmms
  where
  cmms :: Expr -> [Expr]
cmms (Expr
eo :$ Expr
ex :$ Expr
ey)  |  Expr -> Bool
isValue Expr
eo Bool -> Bool -> Bool
&& Expr
eo Expr -> [Expr] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Expr]
cos
                        =  [[Expr]] -> [Expr]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat [ [Expr
eo Expr -> Expr -> Expr
:$ Expr
ex' Expr -> Expr -> Expr
:$ Expr
ey', Expr
eo Expr -> Expr -> Expr
:$ Expr
ey' Expr -> Expr -> Expr
:$ Expr
ex']
                                  | Expr
ex' <- Expr -> [Expr]
cmms Expr
ex
                                  , Expr
ey' <- Expr -> [Expr]
cmms Expr
ey
                                  ]
  cmms (Expr
ef :$ Expr
ex)  =  [ Expr
ef' Expr -> Expr -> Expr
:$ Expr
ex'
                     | Expr
ef' <- Expr -> [Expr]
cmms Expr
ef
                     , Expr
ex' <- Expr -> [Expr]
cmms Expr
ex
                     ]
  cmms Expr
e  =  [Expr
e]