-- |
-- Module      : Data.Express.Canon
-- Copyright   : (c) 2019-2021 Rudy Matela
-- License     : 3-Clause BSD  (see the file LICENSE)
-- Maintainer  : Rudy Matela <rudy@matela.com.br>
--
-- Utilities for canonicalizing 'Expr's with variables.
module Data.Express.Canon
  ( canonicalize
  , canonicalizeWith
  , canonicalization
  , canonicalizationWith
  , isCanonical
  , isCanonicalWith
  , canonicalVariations
  , mostGeneralCanonicalVariation
  , mostSpecificCanonicalVariation
  , fastCanonicalVariations
  , fastMostGeneralVariation
  , fastMostSpecificVariation
  )
where

import Data.Express.Basic
import Data.Express.Name
import Data.Express.Instances
import Data.List ((\\))

-- |
-- Like 'canonicalize' but allows customization
-- of the list of variable names.
-- (cf. 'lookupNames', 'variableNamesFromTemplate')
--
-- > > canonicalizeWith (const ["i","j","k","l",...]) (xx -+- yy)
-- > i + j :: Int
--
-- The argument 'Expr' of the argument function allows
-- to provide a different list of names for different types:
--
-- > > let namesFor e
-- > >   | typ e == typeOf (undefined::Char) = variableNamesFromTemplate "c1"
-- > >   | typ e == typeOf (undefined::Int)  = variableNamesFromTemplate "i"
-- > >   | otherwise                         = variableNamesFromTemplate "x"
--
-- > > canonicalizeWith namesFor ((xx -+- ord' dd) -+- (ord' cc -+- yy))
-- > (i + ord c1) + (ord c2 + j) :: Int
canonicalizeWith :: (Expr -> [String]) -> Expr -> Expr
canonicalizeWith :: (Expr -> [String]) -> Expr -> Expr
canonicalizeWith Expr -> [String]
namesFor Expr
e = Expr
e Expr -> [(Expr, Expr)] -> Expr
//- (Expr -> [String]) -> Expr -> [(Expr, Expr)]
canonicalizationWith Expr -> [String]
namesFor Expr
e

-- |
-- Like 'canonicalization' but allows customization
-- of the list of variable names.
-- (cf. 'lookupNames', 'variableNamesFromTemplate')
canonicalizationWith :: (Expr -> [String]) -> Expr -> [(Expr,Expr)]
canonicalizationWith :: (Expr -> [String]) -> Expr -> [(Expr, Expr)]
canonicalizationWith Expr -> [String]
namesFor Expr
e = [Expr] -> [(Expr, Expr)] -> [(Expr, Expr)]
cr (Expr -> [Expr]
vars Expr
e) []
  where
  cr :: [Expr] -> [(Expr,Expr)] -> [(Expr,Expr)]
  cr :: [Expr] -> [(Expr, Expr)] -> [(Expr, Expr)]
cr []     [(Expr, Expr)]
bs  =  [(Expr, Expr)]
bs
  cr (Expr
e:[Expr]
es) [(Expr, Expr)]
bs  =  [Expr] -> [(Expr, Expr)] -> [(Expr, Expr)]
cr [Expr]
es
                ([(Expr, Expr)] -> [(Expr, Expr)])
-> [(Expr, Expr)] -> [(Expr, Expr)]
forall a b. (a -> b) -> a -> b
$  if Expr
e Expr -> [Expr] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` ((Expr, Expr) -> Expr) -> [(Expr, Expr)] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map (Expr, Expr) -> Expr
forall a b. (a, b) -> a
fst [(Expr, Expr)]
bs
                   then [(Expr, Expr)]
bs
                   else (Expr
e, String
n String -> Expr -> Expr
`varAsTypeOf` Expr
e)(Expr, Expr) -> [(Expr, Expr)] -> [(Expr, Expr)]
forall a. a -> [a] -> [a]
:[(Expr, Expr)]
bs
    where
    existingNames :: [String]
existingNames = [String
n | (Expr
_,Value (Char
'_':String
n) Dynamic
_) <- [(Expr, Expr)]
bs]
    freshNames :: [String]
freshNames = Expr -> [String]
namesFor Expr
e [String] -> [String] -> [String]
forall a. Eq a => [a] -> [a] -> [a]
\\ [String]
existingNames
    n :: String
n = [String] -> String
forall a. [a] -> a
head [String]
freshNames

-- |
-- Like 'isCanonical' but allows specifying
-- the list of variable names.
isCanonicalWith :: (Expr -> [String]) -> Expr -> Bool
isCanonicalWith :: (Expr -> [String]) -> Expr -> Bool
isCanonicalWith Expr -> [String]
ti Expr
e = (Expr -> [String]) -> Expr -> Expr
canonicalizeWith Expr -> [String]
ti Expr
e Expr -> Expr -> Bool
forall a. Eq a => a -> a -> Bool
== Expr
e

-- |
-- Canonicalizes an 'Expr' so that variable names appear in order.
-- Variable names are taken from the 'preludeNameInstances'.
--
-- > > canonicalize (xx -+- yy)
-- > x + y :: Int
--
-- > > canonicalize (yy -+- xx)
-- > x + y :: Int
--
-- > > canonicalize (xx -+- xx)
-- > x + x :: Int
--
-- > > canonicalize (yy -+- yy)
-- > x + x :: Int
--
-- Constants are untouched:
--
-- > > canonicalize (jj -+- (zero -+- abs' ii))
-- > x + (y + abs y) :: Int
--
-- This also works for variable functions:
--
-- > > canonicalize (gg yy -+- ff xx -+- gg xx)
-- > (f x + g y) + f y :: Int
canonicalize :: Expr -> Expr
canonicalize :: Expr -> Expr
canonicalize = (Expr -> [String]) -> Expr -> Expr
canonicalizeWith Expr -> [String]
names'

-- |
-- Return a canonicalization of an 'Expr'
-- that makes variable names appear in order
-- using 'names' as provided by 'preludeNameInstances'.
-- By using '//-' it can 'canonicalize' 'Expr's.
--
-- > > canonicalization (gg yy -+- ff xx -+- gg xx)
-- > [ (x :: Int,        y :: Int)
-- > , (f :: Int -> Int, g :: Int -> Int)
-- > , (y :: Int,        x :: Int)
-- > , (g :: Int -> Int, f :: Int -> Int) ]
--
-- > > canonicalization (yy -+- xx -+- yy)
-- > [ (x :: Int, y :: Int)
-- > , (y :: Int, x :: Int) ]
canonicalization :: Expr -> [(Expr,Expr)]
canonicalization :: Expr -> [(Expr, Expr)]
canonicalization = (Expr -> [String]) -> Expr -> [(Expr, Expr)]
canonicalizationWith Expr -> [String]
names'

-- |
-- Returns whether an 'Expr' is canonical:
-- if applying 'canonicalize' is an identity
-- using 'names' as provided by 'preludeNameInstances'.
isCanonical :: Expr -> Bool
isCanonical :: Expr -> Bool
isCanonical = (Expr -> [String]) -> Expr -> Bool
isCanonicalWith Expr -> [String]
names'

-- 'names' lifted over the 'Expr' type for a handful of prelude Name instances.
names' :: Expr -> [String]
names' :: Expr -> [String]
names' = [Expr] -> Expr -> [String]
lookupNames [Expr]
preludeNameInstances

-- |
-- Returns all canonical variations of an 'Expr'
-- by filling holes with variables.
-- Where possible, variations are listed
-- from most general to least general.
--
-- > > canonicalVariations $ i_
-- > [x :: Int]
--
-- > > canonicalVariations $ i_ -+- i_
-- > [ x + y :: Int
-- > , x + x :: Int ]
--
-- > > canonicalVariations $ i_ -+- i_ -+- i_
-- > [ (x + y) + z :: Int
-- > , (x + y) + x :: Int
-- > , (x + y) + y :: Int
-- > , (x + x) + y :: Int
-- > , (x + x) + x :: Int ]
--
-- > > canonicalVariations $ i_ -+- ord' c_
-- > [x + ord c :: Int]
--
-- > > canonicalVariations $ i_ -+- i_ -+- ord' c_
-- > [ (x + y) + ord c :: Int
-- > , (x + x) + ord c :: Int ]
--
-- > > canonicalVariations $ i_ -+- i_ -+- length' (c_ -:- unit c_)
-- > [ (x + y) + length (c:d:"") :: Int
-- > , (x + y) + length (c:c:"") :: Int
-- > , (x + x) + length (c:d:"") :: Int
-- > , (x + x) + length (c:c:"") :: Int ]
--
-- In an expression without holes this functions just returns a singleton list
-- with the expression itself:
--
-- > > canonicalVariations $ val (0 :: Int)
-- > [0 :: Int]
--
-- > > canonicalVariations $ ord' bee
-- > [ord 'b' :: Int]
--
-- When applying this to expressions already containing variables
-- clashes are avoided and these variables are not touched:
--
-- > > canonicalVariations $ i_ -+- ii -+- jj -+- i_
-- > [ x + i + j + y :: Int
-- > , x + i + j + y :: Int ]
--
-- > > canonicalVariations $ ii -+- jj
-- > [i + j :: Int]
--
-- > > canonicalVariations $ xx -+- i_ -+- i_ -+- length' (c_ -:- unit c_) -+- yy
-- > [ (((x + z) + x') + length (c:d:"")) + y :: Int
-- > , (((x + z) + x') + length (c:c:"")) + y :: Int
-- > , (((x + z) + z) + length (c:d:"")) + y :: Int
-- > , (((x + z) + z) + length (c:c:"")) + y :: Int
-- > ]
canonicalVariations :: Expr -> [Expr]
canonicalVariations :: Expr -> [Expr]
canonicalVariations Expr
e  =  (Expr -> Expr) -> [Expr] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map ([Expr] -> Expr -> Expr
canonicalizeKeeping (Expr -> [Expr]
nonHoleVars Expr
e))
                       ([Expr] -> [Expr]) -> [Expr] -> [Expr]
forall a b. (a -> b) -> a -> b
$  Expr -> [Expr]
fastCanonicalVariations Expr
e

-- |
-- Returns the most general canonical variation of an 'Expr'
-- by filling holes with variables.
--
-- > > mostGeneralCanonicalVariation $ i_
-- > x :: Int
--
-- > > mostGeneralCanonicalVariation $ i_ -+- i_
-- > x + y :: Int
--
-- > > mostGeneralCanonicalVariation $ i_ -+- i_ -+- i_
-- > (x + y) + z :: Int
--
-- > > mostGeneralCanonicalVariation $ i_ -+- ord' c_
-- > x + ord c :: Int
--
-- > > mostGeneralCanonicalVariation $ i_ -+- i_ -+- ord' c_
-- > (x + y) + ord c :: Int
--
-- > > mostGeneralCanonicalVariation $ i_ -+- i_ -+- length' (c_ -:- unit c_)
-- > (x + y) + length (c:d:"") :: Int
--
-- In an expression without holes this functions just returns
-- the given expression itself:
--
-- > > mostGeneralCanonicalVariation $ val (0 :: Int)
-- > 0 :: Int
--
-- > > mostGeneralCanonicalVariation $ ord' bee
-- > ord 'b' :: Int
--
-- This function is the same as taking the 'head' of 'canonicalVariations'
-- but a bit faster.
mostGeneralCanonicalVariation :: Expr -> Expr
mostGeneralCanonicalVariation :: Expr -> Expr
mostGeneralCanonicalVariation Expr
e  =  [Expr] -> Expr -> Expr
canonicalizeKeeping (Expr -> [Expr]
nonHoleVars Expr
e)
                                 (Expr -> Expr) -> Expr -> Expr
forall a b. (a -> b) -> a -> b
$  Expr -> Expr
fastMostGeneralVariation Expr
e

-- |
-- Returns the most specific canonical variation of an 'Expr'
-- by filling holes with variables.
--
-- > > mostSpecificCanonicalVariation $ i_
-- > x :: Int
--
-- > > mostSpecificCanonicalVariation $ i_ -+- i_
-- > x + x :: Int
--
-- > > mostSpecificCanonicalVariation $ i_ -+- i_ -+- i_
-- > (x + x) + x :: Int
--
-- > > mostSpecificCanonicalVariation $ i_ -+- ord' c_
-- > x + ord c :: Int
--
-- > > mostSpecificCanonicalVariation $ i_ -+- i_ -+- ord' c_
-- > (x + x) + ord c :: Int
--
-- > > mostSpecificCanonicalVariation $ i_ -+- i_ -+- length' (c_ -:- unit c_)
-- > (x + x) + length (c:c:"") :: Int
--
-- In an expression without holes this functions just returns
-- the given expression itself:
--
-- > > mostSpecificCanonicalVariation $ val (0 :: Int)
-- > 0 :: Int
--
-- > > mostSpecificCanonicalVariation $ ord' bee
-- > ord 'b' :: Int
--
-- This function is the same as taking the 'last' of 'canonicalVariations'
-- but a bit faster.
mostSpecificCanonicalVariation :: Expr -> Expr
mostSpecificCanonicalVariation :: Expr -> Expr
mostSpecificCanonicalVariation Expr
e  =  [Expr] -> Expr -> Expr
canonicalizeKeeping (Expr -> [Expr]
nonHoleVars Expr
e)
                                  (Expr -> Expr) -> Expr -> Expr
forall a b. (a -> b) -> a -> b
$  Expr -> Expr
fastMostSpecificVariation Expr
e

-- |
-- A faster version of 'canonicalVariations' that
-- disregards name clashes across different types.
-- Results are confusing to the user
-- but fine for Express which differentiates
-- between variables with the same name but different types.
--
-- Without applying 'canonicalize', the following 'Expr'
-- may seem to have only one variable:
--
-- > > fastCanonicalVariations $ i_ -+- ord' c_
-- > [x + ord x :: Int]
--
-- Where in fact it has two, as the second @ x @ has a different type.
-- Applying 'canonicalize' disambiguates:
--
-- > > map canonicalize . fastCanonicalVariations $ i_ -+- ord' c_
-- > [x + ord c :: Int]
--
-- This function is useful when resulting 'Expr's are
-- not intended to be presented to the user
-- but instead to be used by another function.
-- It is simply faster to skip the step where clashes are resolved.
fastCanonicalVariations :: Expr -> [Expr]
fastCanonicalVariations :: Expr -> [Expr]
fastCanonicalVariations Expr
e
  | [Expr] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Expr]
hs'   =  [Expr
e]
  | Bool
otherwise  =  (Expr -> [Expr]) -> [Expr] -> [Expr]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Expr -> [Expr]
fastCanonicalVariations
               ([Expr] -> [Expr]) -> ([Expr] -> [Expr]) -> [Expr] -> [Expr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
.  ([Expr] -> Expr) -> [[Expr]] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map (Expr -> [Expr] -> Expr
fill Expr
e) ([[Expr]] -> [Expr]) -> ([Expr] -> [[Expr]]) -> [Expr] -> [Expr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Expr] -> [[Expr]]
fillings Int
0
               ([Expr] -> [Expr]) -> [Expr] -> [Expr]
forall a b. (a -> b) -> a -> b
$  [Expr
h | Expr
h <- [Expr]
hs', Expr -> TypeRep
typ Expr
h TypeRep -> TypeRep -> Bool
forall a. Eq a => a -> a -> Bool
== Expr -> TypeRep
typ Expr
h']
  where
  hs' :: [Expr]
hs' = Expr -> [Expr]
holes Expr
e
  h' :: Expr
h' = [Expr] -> Expr
forall a. [a] -> a
head [Expr]
hs'

  names :: [String]
names  =  String -> [String]
variableNamesFromTemplate String
"x" [String] -> [String] -> [String]
forall a. Eq a => [a] -> [a] -> [a]
\\ Expr -> [String]
varnames Expr
e

  fillings :: Int -> [Expr] -> [[Expr]]
  fillings :: Int -> [Expr] -> [[Expr]]
fillings Int
i []  =  [[]] -- no holes, single empty filling
  fillings Int
i (Expr
h:[Expr]
hs)  =
    [[[Expr]]] -> [[Expr]]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[[Expr]]] -> [[Expr]]) -> [[[Expr]]] -> [[Expr]]
forall a b. (a -> b) -> a -> b
$ ([Expr] -> [Expr]) -> [[Expr]] -> [[Expr]]
forall a b. (a -> b) -> [a] -> [b]
map ([String]
names [String] -> Int -> String
forall a. [a] -> Int -> a
!! Int
i String -> Expr -> Expr
`varAsTypeOf` Expr
hExpr -> [Expr] -> [Expr]
forall a. a -> [a] -> [a]
:) (Int -> [Expr] -> [[Expr]]
fillings (Int
iInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1) [Expr]
hs) -- new var
           [[Expr]] -> [[[Expr]]] -> [[[Expr]]]
forall a. a -> [a] -> [a]
: [ ([Expr] -> [Expr]) -> [[Expr]] -> [[Expr]]
forall a b. (a -> b) -> [a] -> [b]
map (String
n String -> Expr -> Expr
`varAsTypeOf` Expr
hExpr -> [Expr] -> [Expr]
forall a. a -> [a] -> [a]
:) (Int -> [Expr] -> [[Expr]]
fillings Int
i [Expr]
hs) -- no new variable
             | String
n <- Int -> [String] -> [String]
forall a. Int -> [a] -> [a]
take Int
i [String]
names ]

-- |
-- A faster version of 'mostGeneralCanonicalVariation'
-- that disregards name clashes across different types.
-- Consider using 'mostGeneralCanonicalVariation' instead.
--
-- The same caveats of 'fastCanonicalVariations' do apply here.
fastMostGeneralVariation :: Expr -> Expr
fastMostGeneralVariation :: Expr -> Expr
fastMostGeneralVariation Expr
e  =  Expr -> [Expr] -> Expr
fill Expr
e ((String -> Expr -> Expr) -> [String] -> [Expr] -> [Expr]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith String -> Expr -> Expr
varAsTypeOf [String]
names (Expr -> [Expr]
holes Expr
e))
  where
  names :: [String]
names  =  String -> [String]
variableNamesFromTemplate String
"x" [String] -> [String] -> [String]
forall a. Eq a => [a] -> [a] -> [a]
\\ Expr -> [String]
varnames Expr
e

-- |
-- A faster version of 'mostSpecificCanonicalVariation'
-- that disregards name clashes across different types.
-- Consider using 'mostSpecificCanonicalVariation' instead.
--
-- The same caveats of 'fastCanonicalVariations' do apply here.
fastMostSpecificVariation :: Expr -> Expr
fastMostSpecificVariation :: Expr -> Expr
fastMostSpecificVariation Expr
e  =  Expr -> [Expr] -> Expr
fill Expr
e ((Expr -> Expr) -> [Expr] -> [Expr]
forall a b. (a -> b) -> [a] -> [b]
map (String
name String -> Expr -> Expr
`varAsTypeOf`) (Expr -> [Expr]
holes Expr
e))
  where
  name :: String
name  =  [String] -> String
forall a. [a] -> a
head ([String] -> String) -> [String] -> String
forall a b. (a -> b) -> a -> b
$ String -> [String]
variableNamesFromTemplate String
"x" [String] -> [String] -> [String]
forall a. Eq a => [a] -> [a] -> [a]
\\ Expr -> [String]
varnames Expr
e

-- |
-- Variable names existing in a given Expr.
--
-- This function is not exported.
varnames :: Expr -> [String]
varnames :: Expr -> [String]
varnames Expr
e  =  [String
n | Value (Char
'_':String
n) Dynamic
_ <- Expr -> [Expr]
vars Expr
e]

-- |
-- Variables that are not holes.
--
-- This function is not exported.
nonHoleVars :: Expr -> [Expr]
nonHoleVars :: Expr -> [Expr]
nonHoleVars  =  (Expr -> Bool) -> [Expr] -> [Expr]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Expr -> Bool) -> Expr -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Bool
isHole) ([Expr] -> [Expr]) -> (Expr -> [Expr]) -> Expr -> [Expr]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> [Expr]
nubVars

-- | Canonicalizes an 'Expr' while keeping the given variables untouched.
--
-- > > canonicalizeKeeping [zz] (zz -+- ii -+- jj)
-- > z + x + y :: Int
--
-- > > canonicalizeKeeping [ii,jj] (zz -+- ii -+- jj)
-- > x + i + j :: Int
--
-- This function is not exported.
canonicalizeKeeping :: [Expr] -> Expr -> Expr
canonicalizeKeeping :: [Expr] -> Expr -> Expr
canonicalizeKeeping [Expr]
vs Expr
e  =  (Expr -> [String]) -> Expr -> Expr
canonicalizeWith Expr -> [String]
namesFor Expr
e
  where
  nm :: Expr -> String
nm (Value (Char
'_':String
n) Dynamic
_)  =  String
n
  namesFor :: Expr -> [String]
namesFor Expr
v | Expr
v Expr -> [Expr] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Expr]
vs  =  Expr -> String
nm Expr
v String -> [String] -> [String]
forall a. a -> [a] -> [a]
: [String]
forall a. a
err
             | Bool
otherwise    =  Expr -> [String]
names' Expr
v [String] -> [String] -> [String]
forall a. Eq a => [a] -> [a] -> [a]
\\ (Expr -> String) -> [Expr] -> [String]
forall a b. (a -> b) -> [a] -> [b]
map Expr -> String
nm [Expr]
vs
  err :: a
err  =  String -> a
forall a. HasCallStack => String -> a
error String
"Data.Express.canonicalizeKeeping: the impossible happened.  This is definitely a bug."