module Agda.Compiler.Treeless.Simplify (simplifyTTerm) where

import Control.Arrow (second, (***))
import Control.Monad.Reader
import Data.Traversable (traverse)
import qualified Data.List as List

import Agda.Syntax.Treeless
import Agda.Syntax.Internal (Substitution'(..))
import Agda.Syntax.Literal

import Agda.TypeChecking.Monad
import Agda.TypeChecking.Monad.Builtin
import Agda.TypeChecking.Primitive
import Agda.TypeChecking.Substitute

import Agda.Compiler.Treeless.Compare

import Agda.Utils.List
import Agda.Utils.Maybe

import Agda.Utils.Impossible

data SEnv = SEnv
  { SEnv -> Substitution' TTerm
envSubst   :: Substitution' TTerm
  , SEnv -> [(TTerm, TTerm)]
envRewrite :: [(TTerm, TTerm)] }

type S = Reader SEnv

runS :: S a -> a
runS :: S a -> a
runS m :: S a
m = S a -> SEnv -> a
forall r a. Reader r a -> r -> a
runReader S a
m (SEnv -> a) -> SEnv -> a
forall a b. (a -> b) -> a -> b
$ Substitution' TTerm -> [(TTerm, TTerm)] -> SEnv
SEnv Substitution' TTerm
forall a. Substitution' a
IdS []

lookupVar :: Int -> S TTerm
lookupVar :: Int -> S TTerm
lookupVar i :: Int
i = (SEnv -> TTerm) -> S TTerm
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks ((SEnv -> TTerm) -> S TTerm) -> (SEnv -> TTerm) -> S TTerm
forall a b. (a -> b) -> a -> b
$ (Substitution' TTerm -> Int -> TTerm
forall a. Subst a a => Substitution' a -> Int -> a
`lookupS` Int
i) (Substitution' TTerm -> TTerm)
-> (SEnv -> Substitution' TTerm) -> SEnv -> TTerm
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SEnv -> Substitution' TTerm
envSubst

onSubst :: (Substitution' TTerm -> Substitution' TTerm) -> S a -> S a
onSubst :: (Substitution' TTerm -> Substitution' TTerm) -> S a -> S a
onSubst f :: Substitution' TTerm -> Substitution' TTerm
f = (SEnv -> SEnv) -> S a -> S a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((SEnv -> SEnv) -> S a -> S a) -> (SEnv -> SEnv) -> S a -> S a
forall a b. (a -> b) -> a -> b
$ \ env :: SEnv
env -> SEnv
env { envSubst :: Substitution' TTerm
envSubst = Substitution' TTerm -> Substitution' TTerm
f (SEnv -> Substitution' TTerm
envSubst SEnv
env) }

onRewrite :: Substitution' TTerm -> S a -> S a
onRewrite :: Substitution' TTerm -> S a -> S a
onRewrite rho :: Substitution' TTerm
rho = (SEnv -> SEnv) -> S a -> S a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((SEnv -> SEnv) -> S a -> S a) -> (SEnv -> SEnv) -> S a -> S a
forall a b. (a -> b) -> a -> b
$ \ env :: SEnv
env -> SEnv
env { envRewrite :: [(TTerm, TTerm)]
envRewrite = ((TTerm, TTerm) -> (TTerm, TTerm))
-> [(TTerm, TTerm)] -> [(TTerm, TTerm)]
forall a b. (a -> b) -> [a] -> [b]
map (Substitution' TTerm -> TTerm -> TTerm
forall t a. Subst t a => Substitution' t -> a -> a
applySubst Substitution' TTerm
rho (TTerm -> TTerm)
-> (TTerm -> TTerm) -> (TTerm, TTerm) -> (TTerm, TTerm)
forall (a :: * -> * -> *) b c b' c'.
Arrow a =>
a b c -> a b' c' -> a (b, b') (c, c')
*** Substitution' TTerm -> TTerm -> TTerm
forall t a. Subst t a => Substitution' t -> a -> a
applySubst Substitution' TTerm
rho) (SEnv -> [(TTerm, TTerm)]
envRewrite SEnv
env) }

addRewrite :: TTerm -> TTerm -> S a -> S a
addRewrite :: TTerm -> TTerm -> S a -> S a
addRewrite lhs :: TTerm
lhs rhs :: TTerm
rhs = (SEnv -> SEnv) -> S a -> S a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((SEnv -> SEnv) -> S a -> S a) -> (SEnv -> SEnv) -> S a -> S a
forall a b. (a -> b) -> a -> b
$ \ env :: SEnv
env -> SEnv
env { envRewrite :: [(TTerm, TTerm)]
envRewrite = (TTerm
lhs, TTerm
rhs) (TTerm, TTerm) -> [(TTerm, TTerm)] -> [(TTerm, TTerm)]
forall a. a -> [a] -> [a]
: SEnv -> [(TTerm, TTerm)]
envRewrite SEnv
env }

underLams :: Int -> S a -> S a
underLams :: Int -> S a -> S a
underLams i :: Int
i = Substitution' TTerm -> S a -> S a
forall a. Substitution' TTerm -> S a -> S a
onRewrite (Int -> Substitution' TTerm
forall a. Int -> Substitution' a
raiseS Int
i) (S a -> S a) -> (S a -> S a) -> S a -> S a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Substitution' TTerm -> Substitution' TTerm) -> S a -> S a
forall a.
(Substitution' TTerm -> Substitution' TTerm) -> S a -> S a
onSubst (Int -> Substitution' TTerm -> Substitution' TTerm
forall a. Int -> Substitution' a -> Substitution' a
liftS Int
i)

underLam :: S a -> S a
underLam :: S a -> S a
underLam = Int -> S a -> S a
forall a. Int -> S a -> S a
underLams 1

underLet :: TTerm -> S a -> S a
underLet :: TTerm -> S a -> S a
underLet u :: TTerm
u = Substitution' TTerm -> S a -> S a
forall a. Substitution' TTerm -> S a -> S a
onRewrite (Int -> Substitution' TTerm
forall a. Int -> Substitution' a
raiseS 1) (S a -> S a) -> (S a -> S a) -> S a -> S a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Substitution' TTerm -> Substitution' TTerm) -> S a -> S a
forall a.
(Substitution' TTerm -> Substitution' TTerm) -> S a -> S a
onSubst (\rho :: Substitution' TTerm
rho -> Int -> Substitution' TTerm -> Substitution' TTerm
forall a. Int -> Substitution' a -> Substitution' a
wkS 1 (Substitution' TTerm -> Substitution' TTerm)
-> Substitution' TTerm -> Substitution' TTerm
forall a b. (a -> b) -> a -> b
$ TTerm
u TTerm -> Substitution' TTerm -> Substitution' TTerm
forall a. a -> Substitution' a -> Substitution' a
:# Substitution' TTerm
rho)

bindVar :: Int -> TTerm -> S a -> S a
bindVar :: Int -> TTerm -> S a -> S a
bindVar x :: Int
x u :: TTerm
u = (Substitution' TTerm -> Substitution' TTerm) -> S a -> S a
forall a.
(Substitution' TTerm -> Substitution' TTerm) -> S a -> S a
onSubst (Int -> TTerm -> Substitution' TTerm
forall a. Subst a a => Int -> a -> Substitution' a
inplaceS Int
x TTerm
u Substitution' TTerm -> Substitution' TTerm -> Substitution' TTerm
forall a.
Subst a a =>
Substitution' a -> Substitution' a -> Substitution' a
`composeS`)

rewrite :: TTerm -> S TTerm
rewrite :: TTerm -> S TTerm
rewrite t :: TTerm
t = do
  [(TTerm, TTerm)]
rules <- (SEnv -> [(TTerm, TTerm)])
-> ReaderT SEnv Identity [(TTerm, TTerm)]
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks SEnv -> [(TTerm, TTerm)]
envRewrite
  case [ TTerm
rhs | (lhs :: TTerm
lhs, rhs :: TTerm
rhs) <- [(TTerm, TTerm)]
rules, TTerm -> TTerm -> Bool
equalTerms TTerm
t TTerm
lhs ] of
    rhs :: TTerm
rhs : _ -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
rhs
    []      -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t

data FunctionKit = FunctionKit
  { FunctionKit -> Maybe QName
modAux, FunctionKit -> Maybe QName
divAux, FunctionKit -> Maybe QName
natMinus, FunctionKit -> Maybe QName
true, FunctionKit -> Maybe QName
false :: Maybe QName }

simplifyTTerm :: TTerm -> TCM TTerm
simplifyTTerm :: TTerm -> TCM TTerm
simplifyTTerm t :: TTerm
t = do
  FunctionKit
kit <- Maybe QName
-> Maybe QName
-> Maybe QName
-> Maybe QName
-> Maybe QName
-> FunctionKit
FunctionKit (Maybe QName
 -> Maybe QName
 -> Maybe QName
 -> Maybe QName
 -> Maybe QName
 -> FunctionKit)
-> TCMT IO (Maybe QName)
-> TCMT
     IO
     (Maybe QName
      -> Maybe QName -> Maybe QName -> Maybe QName -> FunctionKit)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> String -> TCMT IO (Maybe QName)
getBuiltinName String
builtinNatModSucAux
                     TCMT
  IO
  (Maybe QName
   -> Maybe QName -> Maybe QName -> Maybe QName -> FunctionKit)
-> TCMT IO (Maybe QName)
-> TCMT
     IO (Maybe QName -> Maybe QName -> Maybe QName -> FunctionKit)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> TCMT IO (Maybe QName)
getBuiltinName String
builtinNatDivSucAux
                     TCMT IO (Maybe QName -> Maybe QName -> Maybe QName -> FunctionKit)
-> TCMT IO (Maybe QName)
-> TCMT IO (Maybe QName -> Maybe QName -> FunctionKit)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> TCMT IO (Maybe QName)
getBuiltinName String
builtinNatMinus
                     TCMT IO (Maybe QName -> Maybe QName -> FunctionKit)
-> TCMT IO (Maybe QName) -> TCMT IO (Maybe QName -> FunctionKit)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> TCMT IO (Maybe QName)
getBuiltinName String
builtinTrue
                     TCMT IO (Maybe QName -> FunctionKit)
-> TCMT IO (Maybe QName) -> TCMT IO FunctionKit
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> String -> TCMT IO (Maybe QName)
getBuiltinName String
builtinFalse
  TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> TCM TTerm) -> TTerm -> TCM TTerm
forall a b. (a -> b) -> a -> b
$ S TTerm -> TTerm
forall a. S a -> a
runS (S TTerm -> TTerm) -> S TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ FunctionKit -> TTerm -> S TTerm
simplify FunctionKit
kit TTerm
t

simplify :: FunctionKit -> TTerm -> S TTerm
simplify :: FunctionKit -> TTerm -> S TTerm
simplify FunctionKit{..} = TTerm -> S TTerm
simpl
  where
    simpl :: TTerm -> S TTerm
simpl = TTerm -> S TTerm
rewrite' (TTerm -> S TTerm) -> (TTerm -> S TTerm) -> TTerm -> S TTerm
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> TTerm -> S TTerm
unchainCase (TTerm -> S TTerm) -> (TTerm -> S TTerm) -> TTerm -> S TTerm
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> \case

      t :: TTerm
t@TDef{}  -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      t :: TTerm
t@TPrim{} -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      t :: TTerm
t@TVar{}  -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t

      TApp (TDef f :: QName
f) [TLit (LitNat _ 0), m :: TTerm
m, n :: TTerm
n, m' :: TTerm
m']
        -- div/mod are equivalent to quot/rem on natural numbers.
        | TTerm
m TTerm -> TTerm -> Bool
forall a. Eq a => a -> a -> Bool
== TTerm
m', QName -> Maybe QName
forall a. a -> Maybe a
Just QName
f Maybe QName -> Maybe QName -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe QName
divAux -> TTerm -> S TTerm
simpl (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PQuot TTerm
n (Integer -> TTerm -> TTerm
tPlusK 1 TTerm
m)
        | TTerm
m TTerm -> TTerm -> Bool
forall a. Eq a => a -> a -> Bool
== TTerm
m', QName -> Maybe QName
forall a. a -> Maybe a
Just QName
f Maybe QName -> Maybe QName -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe QName
modAux -> TTerm -> S TTerm
simpl (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PRem TTerm
n (Integer -> TTerm -> TTerm
tPlusK 1 TTerm
m)

      -- Word64 primitives --

      --  toWord (a ∙ b) == toWord a ∙64 toWord b
      TPFn PITo64 (TPOp op :: TPrim
op a :: TTerm
a b :: TTerm
b)
        | Just op64 :: TPrim
op64 <- TPrim -> Maybe TPrim
opTo64 TPrim
op -> TTerm -> S TTerm
simpl (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
op64 (TPrim -> TTerm -> TTerm
TPFn TPrim
PITo64 TTerm
a) (TPrim -> TTerm -> TTerm
TPFn TPrim
PITo64 TTerm
b)
        where
          opTo64 :: TPrim -> Maybe TPrim
opTo64 op :: TPrim
op = TPrim -> [(TPrim, TPrim)] -> Maybe TPrim
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup TPrim
op [(TPrim
PAdd, TPrim
PAdd64), (TPrim
PSub, TPrim
PSub64), (TPrim
PMul, TPrim
PMul64),
                                 (TPrim
PQuot, TPrim
PQuot64), (TPrim
PRem, TPrim
PRem64)]

      t :: TTerm
t@(TApp (TPrim _) _) -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t  -- taken care of by rewrite'

      TCoerce t :: TTerm
t -> TTerm -> TTerm
TCoerce (TTerm -> TTerm) -> S TTerm -> S TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm
simpl TTerm
t

      TApp f :: TTerm
f es :: Args
es -> do
        TTerm
f  <- TTerm -> S TTerm
simpl TTerm
f
        Args
es <- (TTerm -> S TTerm) -> Args -> ReaderT SEnv Identity Args
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse TTerm -> S TTerm
simpl Args
es
        TTerm -> Args -> S TTerm
maybeMinusToPrim TTerm
f Args
es
      TLam b :: TTerm
b    -> TTerm -> TTerm
TLam (TTerm -> TTerm) -> S TTerm -> S TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> S TTerm -> S TTerm
forall a. S a -> S a
underLam (TTerm -> S TTerm
simpl TTerm
b)
      t :: TTerm
t@TLit{}  -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      t :: TTerm
t@TCon{}  -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      TLet e :: TTerm
e b :: TTerm
b  -> do
        TTerm -> S TTerm
simpl TTerm
e S TTerm -> (TTerm -> S TTerm) -> S TTerm
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          TPFn P64ToI a :: TTerm
a -> do
            -- Inline calls to P64ToI since these trigger optimisations.
            -- Ideally, the optimisations would trigger anyway, but at the
            -- moment they only do if inlining the entire let looks like a
            -- good idea.
            let rho :: Substitution' TTerm
rho = Int -> TTerm -> Substitution' TTerm
forall a. Subst a a => Int -> a -> Substitution' a
inplaceS 0 (TPrim -> TTerm -> TTerm
TPFn TPrim
P64ToI (Int -> TTerm
TVar 0))
            TTerm -> TTerm -> TTerm
tLet TTerm
a (TTerm -> TTerm) -> S TTerm -> S TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm -> S TTerm
forall a. TTerm -> S a -> S a
underLet TTerm
a (TTerm -> S TTerm
simpl (Substitution' TTerm -> TTerm -> TTerm
forall t a. Subst t a => Substitution' t -> a -> a
applySubst Substitution' TTerm
rho TTerm
b))
          e :: TTerm
e -> TTerm -> TTerm -> TTerm
tLet TTerm
e (TTerm -> TTerm) -> S TTerm -> S TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm -> S TTerm
forall a. TTerm -> S a -> S a
underLet TTerm
e (TTerm -> S TTerm
simpl TTerm
b)

      TCase x :: Int
x t :: CaseInfo
t d :: TTerm
d bs :: [TAlt]
bs -> do
        TTerm
v <- Int -> S TTerm
lookupVar Int
x
        let (lets :: Args
lets, u :: TTerm
u) = TTerm -> (Args, TTerm)
tLetView TTerm
v
        case TTerm
u of                          -- TODO: also for literals
          _ | Just (c :: QName
c, as :: Args
as)     <- TTerm -> Maybe (QName, Args)
conView TTerm
u   -> TTerm -> S TTerm
simpl (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ Args -> QName -> Args -> TTerm -> [TAlt] -> TTerm
forall (t :: * -> *).
Foldable t =>
t TTerm -> QName -> Args -> TTerm -> [TAlt] -> TTerm
matchCon Args
lets QName
c Args
as TTerm
d [TAlt]
bs
            | Just (k :: Integer
k, TVar y :: Int
y) <- TTerm -> Maybe (Integer, TTerm)
plusKView TTerm
u -> TTerm -> S TTerm
simpl (TTerm -> S TTerm) -> ([TAlt] -> TTerm) -> [TAlt] -> S TTerm
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Args -> TTerm -> TTerm
forall (t :: * -> *). Foldable t => t TTerm -> TTerm -> TTerm
mkLets Args
lets (TTerm -> TTerm) -> ([TAlt] -> TTerm) -> [TAlt] -> TTerm
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
y CaseInfo
t TTerm
d ([TAlt] -> S TTerm) -> ReaderT SEnv Identity [TAlt] -> S TTerm
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< (TAlt -> ReaderT SEnv Identity TAlt)
-> [TAlt] -> ReaderT SEnv Identity [TAlt]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (Int -> Int -> Integer -> TAlt -> ReaderT SEnv Identity TAlt
matchPlusK Int
y Int
x Integer
k) [TAlt]
bs
          TCase y :: Int
y t1 :: CaseInfo
t1 d1 :: TTerm
d1 bs1 :: [TAlt]
bs1 -> TTerm -> S TTerm
simpl (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ Args -> TTerm -> TTerm
forall (t :: * -> *). Foldable t => t TTerm -> TTerm -> TTerm
mkLets Args
lets (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
y CaseInfo
t1 (TTerm -> TTerm -> TTerm
distrDef TTerm
case1 TTerm
d1) ([TAlt] -> TTerm) -> [TAlt] -> TTerm
forall a b. (a -> b) -> a -> b
$
                                       (TAlt -> TAlt) -> [TAlt] -> [TAlt]
forall a b. (a -> b) -> [a] -> [b]
map (TTerm -> TAlt -> TAlt
distrCase TTerm
case1) [TAlt]
bs1
            where
              -- Γ x Δ -> Γ _ Δ Θ y, where x maps to y and Θ are the lets
              n :: Int
n     = Args -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Args
lets
              rho :: Substitution' TTerm
rho   = Int -> Substitution' TTerm -> Substitution' TTerm
forall a. Int -> Substitution' a -> Substitution' a
liftS (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) (Int -> Substitution' TTerm
forall a. Int -> Substitution' a
raiseS 1)    Substitution' TTerm -> Substitution' TTerm -> Substitution' TTerm
forall a.
Subst a a =>
Substitution' a -> Substitution' a -> Substitution' a
`composeS`
                      Int -> TTerm -> Substitution' TTerm
forall a. DeBruijn a => Int -> a -> Substitution' a
singletonS (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) (Int -> TTerm
TVar 0) Substitution' TTerm -> Substitution' TTerm -> Substitution' TTerm
forall a.
Subst a a =>
Substitution' a -> Substitution' a -> Substitution' a
`composeS`
                      Int -> Substitution' TTerm
forall a. Int -> Substitution' a
raiseS (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1)
              case1 :: TTerm
case1 = Substitution' TTerm -> TTerm -> TTerm
forall t a. Subst t a => Substitution' t -> a -> a
applySubst Substitution' TTerm
rho (Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
x CaseInfo
t TTerm
d [TAlt]
bs)

              distrDef :: TTerm -> TTerm -> TTerm
distrDef v :: TTerm
v d :: TTerm
d | TTerm -> Bool
forall a. Unreachable a => a -> Bool
isUnreachable TTerm
d = TTerm
tUnreachable
                           | Bool
otherwise       = TTerm -> TTerm -> TTerm
tLet TTerm
d TTerm
v

              distrCase :: TTerm -> TAlt -> TAlt
distrCase v :: TTerm
v (TACon c :: QName
c a :: Int
a b :: TTerm
b) = QName -> Int -> TTerm -> TAlt
TACon QName
c Int
a (TTerm -> TAlt) -> TTerm -> TAlt
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm -> TTerm
TLet TTerm
b (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ Int -> Int -> TTerm -> TTerm
forall t a. Subst t a => Int -> Int -> a -> a
raiseFrom 1 Int
a TTerm
v
              distrCase v :: TTerm
v (TALit l :: Literal
l b :: TTerm
b)   = Literal -> TTerm -> TAlt
TALit Literal
l   (TTerm -> TAlt) -> TTerm -> TAlt
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm -> TTerm
TLet TTerm
b TTerm
v
              distrCase v :: TTerm
v (TAGuard g :: TTerm
g b :: TTerm
b) = TTerm -> TTerm -> TAlt
TAGuard TTerm
g (TTerm -> TAlt) -> TTerm -> TAlt
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm -> TTerm
TLet TTerm
b TTerm
v

          _ -> do
            TTerm
d  <- TTerm -> S TTerm
simpl TTerm
d
            [TAlt]
bs <- (TAlt -> ReaderT SEnv Identity TAlt)
-> [TAlt] -> ReaderT SEnv Identity [TAlt]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
traverse (Int -> TAlt -> ReaderT SEnv Identity TAlt
simplAlt Int
x) [TAlt]
bs
            Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase Int
x CaseInfo
t TTerm
d [TAlt]
bs

      t :: TTerm
t@TTerm
TUnit    -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      t :: TTerm
t@TTerm
TSort    -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      t :: TTerm
t@TTerm
TErased  -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      t :: TTerm
t@TError{} -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t

    conView :: TTerm -> Maybe (QName, Args)
conView (TCon c :: QName
c)    = (QName, Args) -> Maybe (QName, Args)
forall a. a -> Maybe a
Just (QName
c, [])
    conView (TApp f :: TTerm
f as :: Args
as) = (Args -> Args) -> (QName, Args) -> (QName, Args)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (Args -> Args -> Args
forall a. [a] -> [a] -> [a]
++ Args
as) ((QName, Args) -> (QName, Args))
-> Maybe (QName, Args) -> Maybe (QName, Args)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> Maybe (QName, Args)
conView TTerm
f
    conView e :: TTerm
e           = Maybe (QName, Args)
forall a. Maybe a
Nothing

    -- Collapse chained cases (case x of bs -> vs; _ -> case x of bs' -> vs'  ==>
    --                         case x of bs -> vs; bs' -> vs')
    unchainCase :: TTerm -> S TTerm
    unchainCase :: TTerm -> S TTerm
unchainCase e :: TTerm
e@(TCase x :: Int
x t :: CaseInfo
t d :: TTerm
d bs :: [TAlt]
bs) = do
      let (lets :: Args
lets, u :: TTerm
u) = TTerm -> (Args, TTerm)
tLetView TTerm
d
          k :: Int
k = Args -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Args
lets
      TTerm -> S TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ case TTerm
u of
        TCase y :: Int
y _ d' :: TTerm
d' bs' :: [TAlt]
bs' | Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
k Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y ->
          Args -> TTerm -> TTerm
forall (t :: * -> *). Foldable t => t TTerm -> TTerm -> TTerm
mkLets Args
lets (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
y CaseInfo
t TTerm
d' ([TAlt] -> TTerm) -> [TAlt] -> TTerm
forall a b. (a -> b) -> a -> b
$ Int -> [TAlt] -> [TAlt]
forall t a. Subst t a => Int -> a -> a
raise Int
k [TAlt]
bs [TAlt] -> [TAlt] -> [TAlt]
forall a. [a] -> [a] -> [a]
++ [TAlt]
bs'
        _ -> TTerm
e
    unchainCase e :: TTerm
e = TTerm -> S TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
e


    mkLets :: t TTerm -> TTerm -> TTerm
mkLets es :: t TTerm
es b :: TTerm
b = (TTerm -> TTerm -> TTerm) -> TTerm -> t TTerm -> TTerm
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TTerm -> TTerm -> TTerm
TLet TTerm
b t TTerm
es

    matchCon :: t TTerm -> QName -> Args -> TTerm -> [TAlt] -> TTerm
matchCon _ _ _ d :: TTerm
d [] = TTerm
d
    matchCon lets :: t TTerm
lets c :: QName
c as :: Args
as d :: TTerm
d (TALit{}   : bs :: [TAlt]
bs) = t TTerm -> QName -> Args -> TTerm -> [TAlt] -> TTerm
matchCon t TTerm
lets QName
c Args
as TTerm
d [TAlt]
bs
    matchCon lets :: t TTerm
lets c :: QName
c as :: Args
as d :: TTerm
d (TAGuard{} : bs :: [TAlt]
bs) = t TTerm -> QName -> Args -> TTerm -> [TAlt] -> TTerm
matchCon t TTerm
lets QName
c Args
as TTerm
d [TAlt]
bs
    matchCon lets :: t TTerm
lets c :: QName
c as :: Args
as d :: TTerm
d (TACon c' :: QName
c' a :: Int
a b :: TTerm
b : bs :: [TAlt]
bs)
      | QName
c QName -> QName -> Bool
forall a. Eq a => a -> a -> Bool
== QName
c'        = (TTerm -> t TTerm -> TTerm) -> t TTerm -> TTerm -> TTerm
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((TTerm -> TTerm -> TTerm) -> TTerm -> t TTerm -> TTerm
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr TTerm -> TTerm -> TTerm
TLet) t TTerm
lets (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ Int -> Args -> TTerm -> TTerm
mkLet 0 Args
as (Int -> Int -> TTerm -> TTerm
forall t a. Subst t a => Int -> Int -> a -> a
raiseFrom Int
a (t TTerm -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t TTerm
lets) TTerm
b)
      | Bool
otherwise      = t TTerm -> QName -> Args -> TTerm -> [TAlt] -> TTerm
matchCon t TTerm
lets QName
c Args
as TTerm
d [TAlt]
bs
      where
        mkLet :: Int -> Args -> TTerm -> TTerm
mkLet _ []       b :: TTerm
b = TTerm
b
        mkLet i :: Int
i (a :: TTerm
a : as :: Args
as) b :: TTerm
b = TTerm -> TTerm -> TTerm
TLet (Int -> TTerm -> TTerm
forall t a. Subst t a => Int -> a -> a
raise Int
i TTerm
a) (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ Int -> Args -> TTerm -> TTerm
mkLet (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ 1) Args
as TTerm
b

    -- Simplify let y = x + k in case y of j     -> u; _ | g[y]     -> v
    -- to       let y = x + k in case x of j - k -> u; _ | g[x + k] -> v
    matchPlusK :: Int -> Int -> Integer -> TAlt -> S TAlt
    matchPlusK :: Int -> Int -> Integer -> TAlt -> ReaderT SEnv Identity TAlt
matchPlusK x :: Int
x y :: Int
y k :: Integer
k (TALit (LitNat r :: Range
r j :: Integer
j) b :: TTerm
b) = TAlt -> ReaderT SEnv Identity TAlt
forall (m :: * -> *) a. Monad m => a -> m a
return (TAlt -> ReaderT SEnv Identity TAlt)
-> TAlt -> ReaderT SEnv Identity TAlt
forall a b. (a -> b) -> a -> b
$ Literal -> TTerm -> TAlt
TALit (Range -> Integer -> Literal
LitNat Range
r (Integer
j Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
k)) TTerm
b
    matchPlusK x :: Int
x y :: Int
y k :: Integer
k (TAGuard g :: TTerm
g b :: TTerm
b) = (TTerm -> TTerm -> TAlt) -> TTerm -> TTerm -> TAlt
forall a b c. (a -> b -> c) -> b -> a -> c
flip TTerm -> TTerm -> TAlt
TAGuard TTerm
b (TTerm -> TAlt) -> S TTerm -> ReaderT SEnv Identity TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm
simpl (Substitution' TTerm -> TTerm -> TTerm
forall t a. Subst t a => Substitution' t -> a -> a
applySubst (Int -> TTerm -> Substitution' TTerm
forall a. Subst a a => Int -> a -> Substitution' a
inplaceS Int
y (Integer -> TTerm -> TTerm
tPlusK Integer
k (Int -> TTerm
TVar Int
x))) TTerm
g)
    matchPlusK x :: Int
x y :: Int
y k :: Integer
k TACon{} = ReaderT SEnv Identity TAlt
forall a. HasCallStack => a
__IMPOSSIBLE__
    matchPlusK x :: Int
x y :: Int
y k :: Integer
k TALit{} = ReaderT SEnv Identity TAlt
forall a. HasCallStack => a
__IMPOSSIBLE__

    simplPrim :: TTerm -> S TTerm
simplPrim (TApp f :: TTerm
f@TPrim{} args :: Args
args) = do
        Args
args    <- (TTerm -> S TTerm) -> Args -> ReaderT SEnv Identity Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> S TTerm
simpl Args
args
        Args
inlined <- (TTerm -> S TTerm) -> Args -> ReaderT SEnv Identity Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> S TTerm
inline Args
args
        let u :: TTerm
u = TTerm -> Args -> TTerm
TApp TTerm
f Args
args
            v :: TTerm
v = TTerm -> TTerm
simplPrim' (TTerm -> Args -> TTerm
TApp TTerm
f Args
inlined)
        TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ if TTerm
v TTerm -> TTerm -> Bool
`betterThan` TTerm
u then TTerm
v else TTerm
u
      where
        inline :: TTerm -> S TTerm
inline (TVar x :: Int
x)                   = do
          TTerm
v <- Int -> S TTerm
lookupVar Int
x
          if TTerm
v TTerm -> TTerm -> Bool
forall a. Eq a => a -> a -> Bool
== Int -> TTerm
TVar Int
x then TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
v else TTerm -> S TTerm
inline TTerm
v
        inline (TApp f :: TTerm
f@TPrim{} args :: Args
args)      = TTerm -> Args -> TTerm
TApp TTerm
f (Args -> TTerm) -> ReaderT SEnv Identity Args -> S TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> S TTerm) -> Args -> ReaderT SEnv Identity Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> S TTerm
inline Args
args
        inline u :: TTerm
u@(TLet _ (TCase 0 _ _ _)) = TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
u
        inline (TLet e :: TTerm
e b :: TTerm
b)                 = TTerm -> S TTerm
inline (Int -> TTerm -> TTerm -> TTerm
forall t a. Subst t a => Int -> t -> a -> a
subst 0 TTerm
e TTerm
b)
        inline u :: TTerm
u                          = TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
u
    simplPrim t :: TTerm
t = TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t

    simplPrim' :: TTerm -> TTerm
    simplPrim' :: TTerm -> TTerm
simplPrim' (TApp (TPrim PSeq) (u :: TTerm
u : v :: TTerm
v : vs :: Args
vs))
      | TTerm
u TTerm -> TTerm -> Bool
forall a. Eq a => a -> a -> Bool
== TTerm
v             = TTerm -> Args -> TTerm
mkTApp TTerm
v Args
vs
      | TApp TCon{} _ <- TTerm
u = TTerm -> Args -> TTerm
mkTApp TTerm
v Args
vs
      | TApp TLit{} _ <- TTerm
u = TTerm -> Args -> TTerm
mkTApp TTerm
v Args
vs
    simplPrim' (TApp (TPrim PLt) [u :: TTerm
u, v :: TTerm
v])
      | Just (PAdd, k :: Integer
k, u :: TTerm
u) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
u,
        Just (PAdd, j :: Integer
j, v :: TTerm
v) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
v,
        Integer
k Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
j = TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PLt TTerm
u TTerm
v
      | Just (PAdd, k :: Integer
k, v :: TTerm
v) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
v,
        TApp (TPrim P64ToI) [u :: TTerm
u] <- TTerm
u,
        Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= 2Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^64, Just trueCon :: QName
trueCon <- Maybe QName
true = QName -> TTerm
TCon QName
trueCon
      | Just k :: Integer
k <- TTerm -> Maybe Integer
intView TTerm
u
      , Just j :: Integer
j <- TTerm -> Maybe Integer
intView TTerm
v
      , Just trueCon :: QName
trueCon <- Maybe QName
true
      , Just falseCon :: QName
falseCon <- Maybe QName
false = if Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
j then QName -> TTerm
TCon QName
trueCon else QName -> TTerm
TCon QName
falseCon
    simplPrim' (TApp (TPrim op :: TPrim
op) [u :: TTerm
u, v :: TTerm
v])
      | TPrim
op TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PGeq, TPrim
PLt, TPrim
PEqI]
      , Just (PAdd, k :: Integer
k, u :: TTerm
u) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
u
      , Just j :: Integer
j <- TTerm -> Maybe Integer
intView TTerm
v = TTerm -> Args -> TTerm
TApp (TPrim -> TTerm
TPrim TPrim
op) [TTerm
u, Integer -> TTerm
tInt (Integer
j Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
k)]
    simplPrim' (TApp (TPrim PEqI) [u :: TTerm
u, v :: TTerm
v])
      | Just (op1 :: TPrim
op1, k :: Integer
k, u :: TTerm
u) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
u,
        Just (op2 :: TPrim
op2, j :: Integer
j, v :: TTerm
v) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
v,
        TPrim
op1 TPrim -> TPrim -> Bool
forall a. Eq a => a -> a -> Bool
== TPrim
op2, Integer
k Integer -> Integer -> Bool
forall a. Eq a => a -> a -> Bool
== Integer
j,
        TPrim
op1 TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PAdd, TPrim
PSub] = TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PEqI TTerm
u TTerm
v
    simplPrim' (TPOp op :: TPrim
op u :: TTerm
u v :: TTerm
v)
      | Bool
zeroL, Bool
isMul Bool -> Bool -> Bool
|| Bool
isDiv = Integer -> TTerm
tInt 0
      | Bool
zeroL, Bool
isAdd          = TTerm
v
      | Bool
zeroR, Bool
isMul          = Integer -> TTerm
tInt 0
      | Bool
zeroR, Bool
isAdd Bool -> Bool -> Bool
|| Bool
isSub = TTerm
u
      where zeroL :: Bool
zeroL = Integer -> Maybe Integer
forall a. a -> Maybe a
Just 0 Maybe Integer -> Maybe Integer -> Bool
forall a. Eq a => a -> a -> Bool
== TTerm -> Maybe Integer
intView TTerm
u Bool -> Bool -> Bool
|| Word64 -> Maybe Word64
forall a. a -> Maybe a
Just 0 Maybe Word64 -> Maybe Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== TTerm -> Maybe Word64
word64View TTerm
u
            zeroR :: Bool
zeroR = Integer -> Maybe Integer
forall a. a -> Maybe a
Just 0 Maybe Integer -> Maybe Integer -> Bool
forall a. Eq a => a -> a -> Bool
== TTerm -> Maybe Integer
intView TTerm
v Bool -> Bool -> Bool
|| Word64 -> Maybe Word64
forall a. a -> Maybe a
Just 0 Maybe Word64 -> Maybe Word64 -> Bool
forall a. Eq a => a -> a -> Bool
== TTerm -> Maybe Word64
word64View TTerm
v
            isAdd :: Bool
isAdd = TPrim
op TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PAdd, TPrim
PAdd64]
            isSub :: Bool
isSub = TPrim
op TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PSub, TPrim
PSub64]
            isMul :: Bool
isMul = TPrim
op TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PMul, TPrim
PMul64]
            isDiv :: Bool
isDiv = TPrim
op TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PQuot, TPrim
PQuot64, TPrim
PRem, TPrim
PRem64]
    simplPrim' (TApp (TPrim op :: TPrim
op) [u :: TTerm
u, v :: TTerm
v])
      | Just u :: TTerm
u <- TTerm -> Maybe TTerm
negView TTerm
u,
        Just v :: TTerm
v <- TTerm -> Maybe TTerm
negView TTerm
v,
        TPrim
op TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PMul, TPrim
PQuot] = TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
op TTerm
u TTerm
v
      | Just u :: TTerm
u <- TTerm -> Maybe TTerm
negView TTerm
u,
        TPrim
op TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PMul, TPrim
PQuot] = TTerm -> TTerm
simplArith (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PSub (Integer -> TTerm
tInt 0) (TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
op TTerm
u TTerm
v)
      | Just v :: TTerm
v <- TTerm -> Maybe TTerm
negView TTerm
v,
        TPrim
op TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PMul, TPrim
PQuot] = TTerm -> TTerm
simplArith (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PSub (Integer -> TTerm
tInt 0) (TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
op TTerm
u TTerm
v)
    simplPrim' (TApp (TPrim PRem) [u :: TTerm
u, v :: TTerm
v])
      | Just u :: TTerm
u <- TTerm -> Maybe TTerm
negView TTerm
u  = TTerm -> TTerm
simplArith (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PSub (Integer -> TTerm
tInt 0) (TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PRem TTerm
u (TTerm -> TTerm
unNeg TTerm
v))
      | Just v :: TTerm
v <- TTerm -> Maybe TTerm
negView TTerm
v  = TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PRem TTerm
u TTerm
v

      -- (fromWord a == fromWord b) = (a ==64 b)
    simplPrim' (TPOp op :: TPrim
op (TPFn P64ToI a :: TTerm
a) (TPFn P64ToI b :: TTerm
b))
        | Just op64 :: TPrim
op64 <- TPrim -> Maybe TPrim
opTo64 TPrim
op = TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
op64 TTerm
a TTerm
b
        where
          opTo64 :: TPrim -> Maybe TPrim
opTo64 op :: TPrim
op = TPrim -> [(TPrim, TPrim)] -> Maybe TPrim
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup TPrim
op [(TPrim
PEqI, TPrim
PEq64), (TPrim
PLt, TPrim
PLt64)]

      -- toWord/fromWord k == fromIntegral k
    simplPrim' (TPFn PITo64 (TLit (LitNat r :: Range
r n :: Integer
n)))    = Literal -> TTerm
TLit (Range -> Word64 -> Literal
LitWord64 Range
r (Integer -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n))
    simplPrim' (TPFn P64ToI (TLit (LitWord64 r :: Range
r n :: Word64
n))) = Literal -> TTerm
TLit (Range -> Integer -> Literal
LitNat Range
r    (Word64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
n))

      -- toWord (fromWord a) == a
    simplPrim' (TPFn PITo64 (TPFn P64ToI a :: TTerm
a)) = TTerm
a

    simplPrim' (TApp f :: TTerm
f@(TPrim op :: TPrim
op) [u :: TTerm
u, v :: TTerm
v]) = TTerm -> TTerm
simplArith (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> Args -> TTerm
TApp TTerm
f [TTerm -> TTerm
simplPrim' TTerm
u, TTerm -> TTerm
simplPrim' TTerm
v]
    simplPrim' u :: TTerm
u = TTerm
u

    unNeg :: TTerm -> TTerm
unNeg u :: TTerm
u | Just v :: TTerm
v <- TTerm -> Maybe TTerm
negView TTerm
u = TTerm
v
            | Bool
otherwise           = TTerm
u

    negView :: TTerm -> Maybe TTerm
negView (TApp (TPrim PSub) [a :: TTerm
a, b :: TTerm
b])
      | Just 0 <- TTerm -> Maybe Integer
intView TTerm
a = TTerm -> Maybe TTerm
forall a. a -> Maybe a
Just TTerm
b
    negView _ = Maybe TTerm
forall a. Maybe a
Nothing

    -- Count arithmetic operations
    betterThan :: TTerm -> TTerm -> Bool
betterThan u :: TTerm
u v :: TTerm
v = TTerm -> Integer
forall a. Num a => TTerm -> a
operations TTerm
u Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
<= TTerm -> Integer
forall a. Num a => TTerm -> a
operations TTerm
v
      where
        operations :: TTerm -> a
operations (TApp (TPrim _) [a :: TTerm
a, b :: TTerm
b]) = 1 a -> a -> a
forall a. Num a => a -> a -> a
+ TTerm -> a
operations TTerm
a a -> a -> a
forall a. Num a => a -> a -> a
+ TTerm -> a
operations TTerm
b
        operations (TApp (TPrim PSeq) (a :: TTerm
a : _))
          | TTerm -> Bool
notVar TTerm
a                       = 1000000  -- only seq on variables!
        operations (TApp (TPrim _) [a :: TTerm
a])    = 1 a -> a -> a
forall a. Num a => a -> a -> a
+ TTerm -> a
operations TTerm
a
        operations TVar{}                  = 0
        operations TLit{}                  = 0
        operations TCon{}                  = 0
        operations TDef{}                  = 0
        operations _                       = 1000

        notVar :: TTerm -> Bool
notVar TVar{} = Bool
False
        notVar _      = Bool
True

    rewrite' :: TTerm -> S TTerm
rewrite' t :: TTerm
t = TTerm -> S TTerm
rewrite (TTerm -> S TTerm) -> S TTerm -> S TTerm
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TTerm -> S TTerm
simplPrim TTerm
t

    constArithView :: TTerm -> Maybe (TPrim, Integer, TTerm)
    constArithView :: TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView (TApp (TPrim op :: TPrim
op) [TLit (LitNat _ k :: Integer
k), u :: TTerm
u])
      | TPrim
op TPrim -> [TPrim] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PAdd, TPrim
PSub] = (TPrim, Integer, TTerm) -> Maybe (TPrim, Integer, TTerm)
forall a. a -> Maybe a
Just (TPrim
op, Integer
k, TTerm
u)
    constArithView (TApp (TPrim op :: TPrim
op) [u :: TTerm
u, TLit (LitNat _ k :: Integer
k)])
      | TPrim
op TPrim -> TPrim -> Bool
forall a. Eq a => a -> a -> Bool
== TPrim
PAdd = (TPrim, Integer, TTerm) -> Maybe (TPrim, Integer, TTerm)
forall a. a -> Maybe a
Just (TPrim
op, Integer
k, TTerm
u)
      | TPrim
op TPrim -> TPrim -> Bool
forall a. Eq a => a -> a -> Bool
== TPrim
PSub = (TPrim, Integer, TTerm) -> Maybe (TPrim, Integer, TTerm)
forall a. a -> Maybe a
Just (TPrim
PAdd, -Integer
k, TTerm
u)
    constArithView _ = Maybe (TPrim, Integer, TTerm)
forall a. Maybe a
Nothing

    simplAlt :: Int -> TAlt -> ReaderT SEnv Identity TAlt
simplAlt x :: Int
x (TACon c :: QName
c a :: Int
a b :: TTerm
b) = QName -> Int -> TTerm -> TAlt
TACon QName
c Int
a (TTerm -> TAlt) -> S TTerm -> ReaderT SEnv Identity TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> S TTerm -> S TTerm
forall a. Int -> S a -> S a
underLams Int
a (Int -> TTerm -> S TTerm -> S TTerm
forall a. Int -> TTerm -> S a -> S a
maybeAddRewrite (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
a) TTerm
conTerm (S TTerm -> S TTerm) -> S TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> S TTerm
simpl TTerm
b)
      where conTerm :: TTerm
conTerm = TTerm -> Args -> TTerm
mkTApp (QName -> TTerm
TCon QName
c) (Args -> TTerm) -> Args -> TTerm
forall a b. (a -> b) -> a -> b
$ (Int -> TTerm) -> [Int] -> Args
forall a b. (a -> b) -> [a] -> [b]
map Int -> TTerm
TVar ([Int] -> Args) -> [Int] -> Args
forall a b. (a -> b) -> a -> b
$ Int -> [Int]
forall a. Integral a => a -> [a]
downFrom Int
a
    simplAlt x :: Int
x (TALit l :: Literal
l b :: TTerm
b)   = Literal -> TTerm -> TAlt
TALit Literal
l   (TTerm -> TAlt) -> S TTerm -> ReaderT SEnv Identity TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> TTerm -> S TTerm -> S TTerm
forall a. Int -> TTerm -> S a -> S a
maybeAddRewrite Int
x (Literal -> TTerm
TLit Literal
l) (TTerm -> S TTerm
simpl TTerm
b)
    simplAlt x :: Int
x (TAGuard g :: TTerm
g b :: TTerm
b) = TTerm -> TTerm -> TAlt
TAGuard   (TTerm -> TTerm -> TAlt)
-> S TTerm -> ReaderT SEnv Identity (TTerm -> TAlt)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm
simpl TTerm
g ReaderT SEnv Identity (TTerm -> TAlt)
-> S TTerm -> ReaderT SEnv Identity TAlt
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> S TTerm
simpl TTerm
b

    -- If x is already bound we add a rewrite, otherwise we bind x to rhs.
    maybeAddRewrite :: Int -> TTerm -> S a -> S a
maybeAddRewrite x :: Int
x rhs :: TTerm
rhs cont :: S a
cont = do
      TTerm
v <- Int -> S TTerm
lookupVar Int
x
      case TTerm
v of
        TVar y :: Int
y | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y -> Int -> TTerm -> S a -> S a
forall a. Int -> TTerm -> S a -> S a
bindVar Int
x TTerm
rhs (S a -> S a) -> S a -> S a
forall a b. (a -> b) -> a -> b
$ S a
cont
        _ -> TTerm -> TTerm -> S a -> S a
forall a. TTerm -> TTerm -> S a -> S a
addRewrite TTerm
v TTerm
rhs S a
cont

    isTrue :: TTerm -> Bool
isTrue (TCon c :: QName
c) = QName -> Maybe QName
forall a. a -> Maybe a
Just QName
c Maybe QName -> Maybe QName -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe QName
true
    isTrue _        = Bool
False

    isFalse :: TTerm -> Bool
isFalse (TCon c :: QName
c) = QName -> Maybe QName
forall a. a -> Maybe a
Just QName
c Maybe QName -> Maybe QName -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe QName
false
    isFalse _        = Bool
False

    maybeMinusToPrim :: TTerm -> Args -> S TTerm
maybeMinusToPrim f :: TTerm
f@(TDef minus :: QName
minus) es :: Args
es@[a :: TTerm
a, b :: TTerm
b]
      | QName -> Maybe QName
forall a. a -> Maybe a
Just QName
minus Maybe QName -> Maybe QName -> Bool
forall a. Eq a => a -> a -> Bool
== Maybe QName
natMinus = do
      Bool
leq  <- TTerm -> TTerm -> ReaderT SEnv Identity Bool
forall (m :: * -> *).
MonadReader SEnv m =>
TTerm -> TTerm -> m Bool
checkLeq TTerm
b TTerm
a
      if Bool
leq then TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PSub TTerm
a TTerm
b
             else TTerm -> Args -> S TTerm
tApp TTerm
f Args
es

    maybeMinusToPrim f :: TTerm
f es :: Args
es = TTerm -> Args -> S TTerm
tApp TTerm
f Args
es

    tLet :: TTerm -> TTerm -> TTerm
tLet (TVar x :: Int
x) b :: TTerm
b = Int -> TTerm -> TTerm -> TTerm
forall t a. Subst t a => Int -> t -> a -> a
subst 0 (Int -> TTerm
TVar Int
x) TTerm
b
    tLet e :: TTerm
e (TVar 0) = TTerm
e
    tLet e :: TTerm
e b :: TTerm
b        = TTerm -> TTerm -> TTerm
TLet TTerm
e TTerm
b

    tCase :: Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
    tCase :: Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase x :: Int
x t :: CaseInfo
t d :: TTerm
d [] = TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
d
    tCase x :: Int
x t :: CaseInfo
t d :: TTerm
d bs :: [TAlt]
bs
      | TTerm -> Bool
forall a. Unreachable a => a -> Bool
isUnreachable TTerm
d =
        case [TAlt] -> [TAlt]
forall a. [a] -> [a]
reverse [TAlt]
bs' of
          [] -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
d
          TALit _ b :: TTerm
b   : as :: [TAlt]
as  -> Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase Int
x CaseInfo
t TTerm
b ([TAlt] -> [TAlt]
forall a. [a] -> [a]
reverse [TAlt]
as)
          TAGuard _ b :: TTerm
b : as :: [TAlt]
as  -> Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase Int
x CaseInfo
t TTerm
b ([TAlt] -> [TAlt]
forall a. [a] -> [a]
reverse [TAlt]
as)
          TACon c :: QName
c a :: Int
a b :: TTerm
b : _   -> Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase' Int
x CaseInfo
t TTerm
d [TAlt]
bs'
      | Bool
otherwise = do
        TTerm
d' <- TTerm -> S TTerm
lookupIfVar TTerm
d
        case TTerm
d' of
          TCase y :: Int
y _ d :: TTerm
d bs'' :: [TAlt]
bs'' | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y ->
            Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase Int
x CaseInfo
t TTerm
d ([TAlt]
bs' [TAlt] -> [TAlt] -> [TAlt]
forall a. [a] -> [a] -> [a]
++ (TAlt -> Bool) -> [TAlt] -> [TAlt]
forall a. (a -> Bool) -> [a] -> [a]
filter TAlt -> Bool
noOverlap [TAlt]
bs'')
          _ -> Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase' Int
x CaseInfo
t TTerm
d [TAlt]
bs'
      where
        bs' :: [TAlt]
bs' = (TAlt -> Bool) -> [TAlt] -> [TAlt]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (TAlt -> Bool) -> TAlt -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TAlt -> Bool
forall a. Unreachable a => a -> Bool
isUnreachable) [TAlt]
bs

        lookupIfVar :: TTerm -> S TTerm
lookupIfVar (TVar i :: Int
i) = Int -> S TTerm
lookupVar Int
i
        lookupIfVar t :: TTerm
t = TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t

        noOverlap :: TAlt -> Bool
noOverlap b :: TAlt
b = Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ (TAlt -> Bool) -> [TAlt] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (TAlt -> TAlt -> Bool
overlapped TAlt
b) [TAlt]
bs'
        overlapped :: TAlt -> TAlt -> Bool
overlapped (TACon c :: QName
c _ _)  (TACon c' :: QName
c' _ _) = QName
c QName -> QName -> Bool
forall a. Eq a => a -> a -> Bool
== QName
c'
        overlapped (TALit l :: Literal
l _)    (TALit l' :: Literal
l' _)   = Literal
l Literal -> Literal -> Bool
forall a. Eq a => a -> a -> Bool
== Literal
l'
        overlapped _              _              = Bool
False

    -- | Drop unreachable cases for Nat and Int cases.
    pruneLitCases :: Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
    pruneLitCases :: Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
pruneLitCases x :: Int
x t :: CaseInfo
t d :: TTerm
d bs :: [TAlt]
bs | CaseType
CTNat CaseType -> CaseType -> Bool
forall a. Eq a => a -> a -> Bool
== CaseInfo -> CaseType
caseType CaseInfo
t =
      case [TAlt] -> [Integer] -> Maybe Integer -> Maybe [TAlt]
complete [TAlt]
bs [] Maybe Integer
forall a. Maybe a
Nothing of
        Just bs' :: [TAlt]
bs' -> Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase Int
x CaseInfo
t TTerm
tUnreachable [TAlt]
bs'
        Nothing  -> TTerm -> S TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
x CaseInfo
t TTerm
d [TAlt]
bs
      where
        complete :: [TAlt] -> [Integer] -> Maybe Integer -> Maybe [TAlt]
complete bs :: [TAlt]
bs small :: [Integer]
small (Just upper :: Integer
upper)
          | [Integer] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Integer] -> Bool) -> [Integer] -> Bool
forall a b. (a -> b) -> a -> b
$ [0..Integer
upper Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- 1] [Integer] -> [Integer] -> [Integer]
forall a. Eq a => [a] -> [a] -> [a]
List.\\ [Integer]
small = [TAlt] -> Maybe [TAlt]
forall a. a -> Maybe a
Just []
        complete (b :: TAlt
b@(TALit (LitNat _ n :: Integer
n) _) : bs :: [TAlt]
bs) small :: [Integer]
small upper :: Maybe Integer
upper =
          (TAlt
b TAlt -> [TAlt] -> [TAlt]
forall a. a -> [a] -> [a]
:) ([TAlt] -> [TAlt]) -> Maybe [TAlt] -> Maybe [TAlt]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TAlt] -> [Integer] -> Maybe Integer -> Maybe [TAlt]
complete [TAlt]
bs (Integer
n Integer -> [Integer] -> [Integer]
forall a. a -> [a] -> [a]
: [Integer]
small) Maybe Integer
upper
        complete (b :: TAlt
b@(TAGuard (TApp (TPrim PGeq) [TVar y :: Int
y, TLit (LitNat _ j :: Integer
j)]) _) : bs :: [TAlt]
bs) small :: [Integer]
small upper :: Maybe Integer
upper | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y =
          (TAlt
b TAlt -> [TAlt] -> [TAlt]
forall a. a -> [a] -> [a]
:) ([TAlt] -> [TAlt]) -> Maybe [TAlt] -> Maybe [TAlt]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [TAlt] -> [Integer] -> Maybe Integer -> Maybe [TAlt]
complete [TAlt]
bs [Integer]
small (Integer -> Maybe Integer
forall a. a -> Maybe a
Just (Integer -> Maybe Integer) -> Integer -> Maybe Integer
forall a b. (a -> b) -> a -> b
$ Integer -> (Integer -> Integer) -> Maybe Integer -> Integer
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Integer
j (Integer -> Integer -> Integer
forall a. Ord a => a -> a -> a
min Integer
j) Maybe Integer
upper)
        complete _ _ _ = Maybe [TAlt]
forall a. Maybe a
Nothing

    pruneLitCases x :: Int
x t :: CaseInfo
t d :: TTerm
d bs :: [TAlt]
bs
      | CaseType
CTInt CaseType -> CaseType -> Bool
forall a. Eq a => a -> a -> Bool
== CaseInfo -> CaseType
caseType CaseInfo
t = TTerm -> S TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
x CaseInfo
t TTerm
d [TAlt]
bs -- TODO
      | Bool
otherwise           = TTerm -> S TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
x CaseInfo
t TTerm
d [TAlt]
bs

    tCase' :: Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase' x :: Int
x t :: CaseInfo
t d :: TTerm
d [] = TTerm -> S TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
d
    tCase' x :: Int
x t :: CaseInfo
t d :: TTerm
d bs :: [TAlt]
bs = Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
pruneLitCases Int
x CaseInfo
t TTerm
d [TAlt]
bs

    tApp :: TTerm -> [TTerm] -> S TTerm
    tApp :: TTerm -> Args -> S TTerm
tApp (TLet e :: TTerm
e b :: TTerm
b) es :: Args
es = TTerm -> TTerm -> TTerm
TLet TTerm
e (TTerm -> TTerm) -> S TTerm -> S TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm -> S TTerm
forall a. TTerm -> S a -> S a
underLet TTerm
e (TTerm -> Args -> S TTerm
tApp TTerm
b (Int -> Args -> Args
forall t a. Subst t a => Int -> a -> a
raise 1 Args
es))
    tApp (TCase x :: Int
x t :: CaseInfo
t d :: TTerm
d bs :: [TAlt]
bs) es :: Args
es = do
      TTerm
d  <- TTerm -> Args -> S TTerm
tApp TTerm
d Args
es
      [TAlt]
bs <- (TAlt -> ReaderT SEnv Identity TAlt)
-> [TAlt] -> ReaderT SEnv Identity [TAlt]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM (TAlt -> Args -> ReaderT SEnv Identity TAlt
`tAppAlt` Args
es) [TAlt]
bs
      TTerm -> S TTerm
simpl (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
x CaseInfo
t TTerm
d [TAlt]
bs    -- will resimplify branches
    tApp (TVar x :: Int
x) es :: Args
es = do
      TTerm
v <- Int -> S TTerm
lookupVar Int
x
      case TTerm
v of
        _ | TTerm
v TTerm -> TTerm -> Bool
forall a. Eq a => a -> a -> Bool
/= Int -> TTerm
TVar Int
x Bool -> Bool -> Bool
&& TTerm -> Bool
isAtomic TTerm
v -> TTerm -> Args -> S TTerm
tApp TTerm
v Args
es
        TLam{} -> TTerm -> Args -> S TTerm
tApp TTerm
v Args
es   -- could blow up the code
        _      -> TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> Args -> TTerm
mkTApp (Int -> TTerm
TVar Int
x) Args
es
    tApp f :: TTerm
f [] = TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
f
    tApp (TLam b :: TTerm
b) (TVar i :: Int
i : es :: Args
es) = TTerm -> Args -> S TTerm
tApp (Int -> TTerm -> TTerm -> TTerm
forall t a. Subst t a => Int -> t -> a -> a
subst 0 (Int -> TTerm
TVar Int
i) TTerm
b) Args
es
    tApp (TLam b :: TTerm
b) (e :: TTerm
e : es :: Args
es) = TTerm -> Args -> S TTerm
tApp (TTerm -> TTerm -> TTerm
TLet TTerm
e TTerm
b) Args
es
    tApp f :: TTerm
f es :: Args
es = TTerm -> S TTerm
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> Args -> TTerm
TApp TTerm
f Args
es

    tAppAlt :: TAlt -> Args -> ReaderT SEnv Identity TAlt
tAppAlt (TACon c :: QName
c a :: Int
a b :: TTerm
b) es :: Args
es = QName -> Int -> TTerm -> TAlt
TACon QName
c Int
a (TTerm -> TAlt) -> S TTerm -> ReaderT SEnv Identity TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> S TTerm -> S TTerm
forall a. Int -> S a -> S a
underLams Int
a (TTerm -> Args -> S TTerm
tApp TTerm
b (Int -> Args -> Args
forall t a. Subst t a => Int -> a -> a
raise Int
a Args
es))
    tAppAlt (TALit l :: Literal
l b :: TTerm
b) es :: Args
es   = Literal -> TTerm -> TAlt
TALit Literal
l   (TTerm -> TAlt) -> S TTerm -> ReaderT SEnv Identity TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> Args -> S TTerm
tApp TTerm
b Args
es
    tAppAlt (TAGuard g :: TTerm
g b :: TTerm
b) es :: Args
es = TTerm -> TTerm -> TAlt
TAGuard TTerm
g (TTerm -> TAlt) -> S TTerm -> ReaderT SEnv Identity TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> Args -> S TTerm
tApp TTerm
b Args
es

    isAtomic :: TTerm -> Bool
isAtomic v :: TTerm
v = case TTerm
v of
      TVar{}    -> Bool
True
      TCon{}    -> Bool
True
      TPrim{}   -> Bool
True
      TDef{}    -> Bool
True
      TLit{}    -> Bool
True
      TSort{}   -> Bool
True
      TErased{} -> Bool
True
      TError{}  -> Bool
True
      _         -> Bool
False

    checkLeq :: TTerm -> TTerm -> m Bool
checkLeq a :: TTerm
a b :: TTerm
b = do
      Substitution' TTerm
rho  <- (SEnv -> Substitution' TTerm) -> m (Substitution' TTerm)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks SEnv -> Substitution' TTerm
envSubst
      [(TTerm, TTerm)]
rwr  <- (SEnv -> [(TTerm, TTerm)]) -> m [(TTerm, TTerm)]
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks SEnv -> [(TTerm, TTerm)]
envRewrite
      let nf :: TTerm -> Arith
nf = TTerm -> Arith
toArith (TTerm -> Arith) -> (TTerm -> TTerm) -> TTerm -> Arith
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Substitution' TTerm -> TTerm -> TTerm
forall t a. Subst t a => Substitution' t -> a -> a
applySubst Substitution' TTerm
rho
          less :: [(Arith, Arith)]
less = [ (TTerm -> Arith
nf TTerm
a, TTerm -> Arith
nf TTerm
b) | (TPOp PLt a :: TTerm
a b :: TTerm
b, rhs :: TTerm
rhs) <- [(TTerm, TTerm)]
rwr, TTerm -> Bool
isTrue  TTerm
rhs ]
          leq :: [(Arith, Arith)]
leq  = [ (TTerm -> Arith
nf TTerm
b, TTerm -> Arith
nf TTerm
a) | (TPOp PLt a :: TTerm
a b :: TTerm
b, rhs :: TTerm
rhs) <- [(TTerm, TTerm)]
rwr, TTerm -> Bool
isFalse TTerm
rhs ]

          match :: (a, a) -> (a, a) -> Maybe a
match (j :: a
j, as :: a
as) (k :: a
k, bs :: a
bs)
            | a
as a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
bs  = a -> Maybe a
forall a. a -> Maybe a
Just (a
j a -> a -> a
forall a. Num a => a -> a -> a
- a
k)
            | Bool
otherwise = Maybe a
forall a. Maybe a
Nothing

          -- Do we have x ≤ y given x' < y' + d ?
          matchEqn :: a -> (a, a) -> (a, a) -> ((a, a), (a, a)) -> Bool
matchEqn d :: a
d x :: (a, a)
x y :: (a, a)
y (x' :: (a, a)
x', y' :: (a, a)
y') = Maybe () -> Bool
forall a. Maybe a -> Bool
isJust (Maybe () -> Bool) -> Maybe () -> Bool
forall a b. (a -> b) -> a -> b
$ do
            a
k <- (a, a) -> (a, a) -> Maybe a
forall a a. (Eq a, Num a) => (a, a) -> (a, a) -> Maybe a
match (a, a)
x (a, a)
x'     -- x = x' + k
            a
j <- (a, a) -> (a, a) -> Maybe a
forall a a. (Eq a, Num a) => (a, a) -> (a, a) -> Maybe a
match (a, a)
y (a, a)
y'     -- y = y' + j
            Bool -> Maybe ()
forall (f :: * -> *). Alternative f => Bool -> f ()
guard (a
k a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
j a -> a -> a
forall a. Num a => a -> a -> a
+ a
d)  -- x ≤ y if k ≤ j + d

          matchLess :: Arith -> Arith -> (Arith, Arith) -> Bool
matchLess = Integer -> Arith -> Arith -> (Arith, Arith) -> Bool
forall a a a.
(Num a, Ord a, Eq a, Eq a) =>
a -> (a, a) -> (a, a) -> ((a, a), (a, a)) -> Bool
matchEqn 1
          matchLeq :: Arith -> Arith -> (Arith, Arith) -> Bool
matchLeq  = Integer -> Arith -> Arith -> (Arith, Arith) -> Bool
forall a a a.
(Num a, Ord a, Eq a, Eq a) =>
a -> (a, a) -> (a, a) -> ((a, a), (a, a)) -> Bool
matchEqn 0

          literal :: (a, [a]) -> (a, [a]) -> Bool
literal (j :: a
j, []) (k :: a
k, []) = a
j a -> a -> Bool
forall a. Ord a => a -> a -> Bool
<= a
k
          literal _ _ = Bool
False

          -- k + fromWord x ≤ y  if  k + 2^64 - 1 ≤ y
          wordUpperBound :: Arith -> Arith -> Bool
wordUpperBound (k :: Integer
k, [Pos (TApp (TPrim P64ToI) _)]) y :: Arith
y = Arith -> Arith -> Bool
go (Integer
k Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ 2Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^64 Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- 1, []) Arith
y
          wordUpperBound _ _ = Bool
False

          -- x ≤ k + fromWord y  if  x ≤ k
          wordLowerBound :: Arith -> Arith -> Bool
wordLowerBound a :: Arith
a (k :: Integer
k, [Pos (TApp (TPrim P64ToI) _)]) = Arith -> Arith -> Bool
go Arith
a (Integer
k, [])
          wordLowerBound _ _ = Bool
False

          go :: Arith -> Arith -> Bool
go x :: Arith
x y :: Arith
y = [Bool] -> Bool
forall (t :: * -> *). Foldable t => t Bool -> Bool
or
            [ Arith -> Arith -> Bool
forall a a a. Ord a => (a, [a]) -> (a, [a]) -> Bool
literal Arith
x Arith
y
            , Arith -> Arith -> Bool
wordUpperBound Arith
x Arith
y
            , Arith -> Arith -> Bool
wordLowerBound Arith
x Arith
y
            , ((Arith, Arith) -> Bool) -> [(Arith, Arith)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Arith -> Arith -> (Arith, Arith) -> Bool
matchLess Arith
x Arith
y) [(Arith, Arith)]
less
            , ((Arith, Arith) -> Bool) -> [(Arith, Arith)] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Arith -> Arith -> (Arith, Arith) -> Bool
matchLeq Arith
x Arith
y) [(Arith, Arith)]
leq ]

      Bool -> m Bool
forall (m :: * -> *) a. Monad m => a -> m a
return (Bool -> m Bool) -> Bool -> m Bool
forall a b. (a -> b) -> a -> b
$ Arith -> Arith -> Bool
go (TTerm -> Arith
nf TTerm
a) (TTerm -> Arith
nf TTerm
b)

type Arith = (Integer, [Atom])

data Atom = Pos TTerm | Neg TTerm
  deriving (Int -> Atom -> ShowS
[Atom] -> ShowS
Atom -> String
(Int -> Atom -> ShowS)
-> (Atom -> String) -> ([Atom] -> ShowS) -> Show Atom
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Atom] -> ShowS
$cshowList :: [Atom] -> ShowS
show :: Atom -> String
$cshow :: Atom -> String
showsPrec :: Int -> Atom -> ShowS
$cshowsPrec :: Int -> Atom -> ShowS
Show, Atom -> Atom -> Bool
(Atom -> Atom -> Bool) -> (Atom -> Atom -> Bool) -> Eq Atom
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Atom -> Atom -> Bool
$c/= :: Atom -> Atom -> Bool
== :: Atom -> Atom -> Bool
$c== :: Atom -> Atom -> Bool
Eq, Eq Atom
Eq Atom =>
(Atom -> Atom -> Ordering)
-> (Atom -> Atom -> Bool)
-> (Atom -> Atom -> Bool)
-> (Atom -> Atom -> Bool)
-> (Atom -> Atom -> Bool)
-> (Atom -> Atom -> Atom)
-> (Atom -> Atom -> Atom)
-> Ord Atom
Atom -> Atom -> Bool
Atom -> Atom -> Ordering
Atom -> Atom -> Atom
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
min :: Atom -> Atom -> Atom
$cmin :: Atom -> Atom -> Atom
max :: Atom -> Atom -> Atom
$cmax :: Atom -> Atom -> Atom
>= :: Atom -> Atom -> Bool
$c>= :: Atom -> Atom -> Bool
> :: Atom -> Atom -> Bool
$c> :: Atom -> Atom -> Bool
<= :: Atom -> Atom -> Bool
$c<= :: Atom -> Atom -> Bool
< :: Atom -> Atom -> Bool
$c< :: Atom -> Atom -> Bool
compare :: Atom -> Atom -> Ordering
$ccompare :: Atom -> Atom -> Ordering
$cp1Ord :: Eq Atom
Ord)

aNeg :: Atom -> Atom
aNeg :: Atom -> Atom
aNeg (Pos a :: TTerm
a) = TTerm -> Atom
Neg TTerm
a
aNeg (Neg a :: TTerm
a) = TTerm -> Atom
Pos TTerm
a

aCancel :: [Atom] -> [Atom]
aCancel :: [Atom] -> [Atom]
aCancel (a :: Atom
a : as :: [Atom]
as)
  | (Atom -> Atom
aNeg Atom
a) Atom -> [Atom] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Atom]
as = [Atom] -> [Atom]
aCancel (Atom -> [Atom] -> [Atom]
forall a. Eq a => a -> [a] -> [a]
List.delete (Atom -> Atom
aNeg Atom
a) [Atom]
as)
  | Bool
otherwise          = Atom
a Atom -> [Atom] -> [Atom]
forall a. a -> [a] -> [a]
: [Atom] -> [Atom]
aCancel [Atom]
as
aCancel [] = []

sortR :: Ord a => [a] -> [a]
sortR :: [a] -> [a]
sortR = (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
List.sortBy ((a -> a -> Ordering) -> a -> a -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare)

aAdd :: Arith -> Arith -> Arith
aAdd :: Arith -> Arith -> Arith
aAdd (a :: Integer
a, xs :: [Atom]
xs) (b :: Integer
b, ys :: [Atom]
ys) = (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
b, [Atom] -> [Atom]
aCancel ([Atom] -> [Atom]) -> [Atom] -> [Atom]
forall a b. (a -> b) -> a -> b
$ [Atom] -> [Atom]
forall a. Ord a => [a] -> [a]
sortR ([Atom] -> [Atom]) -> [Atom] -> [Atom]
forall a b. (a -> b) -> a -> b
$ [Atom]
xs [Atom] -> [Atom] -> [Atom]
forall a. [a] -> [a] -> [a]
++ [Atom]
ys)

aSub :: Arith -> Arith -> Arith
aSub :: Arith -> Arith -> Arith
aSub (a :: Integer
a, xs :: [Atom]
xs) (b :: Integer
b, ys :: [Atom]
ys) = (Integer
a Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
b, [Atom] -> [Atom]
aCancel ([Atom] -> [Atom]) -> [Atom] -> [Atom]
forall a b. (a -> b) -> a -> b
$ [Atom] -> [Atom]
forall a. Ord a => [a] -> [a]
sortR ([Atom] -> [Atom]) -> [Atom] -> [Atom]
forall a b. (a -> b) -> a -> b
$ [Atom]
xs [Atom] -> [Atom] -> [Atom]
forall a. [a] -> [a] -> [a]
++ (Atom -> Atom) -> [Atom] -> [Atom]
forall a b. (a -> b) -> [a] -> [b]
map Atom -> Atom
aNeg [Atom]
ys)

fromArith :: Arith -> TTerm
fromArith :: Arith -> TTerm
fromArith (n :: Integer
n, []) = Integer -> TTerm
tInt Integer
n
fromArith (0, xs :: [Atom]
xs)
  | (ys :: [Atom]
ys, Pos a :: TTerm
a : zs :: [Atom]
zs) <- (Atom -> Bool) -> [Atom] -> ([Atom], [Atom])
forall a. (a -> Bool) -> [a] -> ([a], [a])
break Atom -> Bool
isPos [Atom]
xs = (TTerm -> Atom -> TTerm) -> TTerm -> [Atom] -> TTerm
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl TTerm -> Atom -> TTerm
addAtom TTerm
a ([Atom]
ys [Atom] -> [Atom] -> [Atom]
forall a. [a] -> [a] -> [a]
++ [Atom]
zs)
fromArith (n :: Integer
n, xs :: [Atom]
xs)
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< 0, (ys :: [Atom]
ys, Pos a :: TTerm
a : zs :: [Atom]
zs) <- (Atom -> Bool) -> [Atom] -> ([Atom], [Atom])
forall a. (a -> Bool) -> [a] -> ([a], [a])
break Atom -> Bool
isPos [Atom]
xs =
    TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PSub ((TTerm -> Atom -> TTerm) -> TTerm -> [Atom] -> TTerm
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl TTerm -> Atom -> TTerm
addAtom TTerm
a ([Atom]
ys [Atom] -> [Atom] -> [Atom]
forall a. [a] -> [a] -> [a]
++ [Atom]
zs)) (Integer -> TTerm
tInt (-Integer
n))
fromArith (n :: Integer
n, xs :: [Atom]
xs) = (TTerm -> Atom -> TTerm) -> TTerm -> [Atom] -> TTerm
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl TTerm -> Atom -> TTerm
addAtom (Integer -> TTerm
tInt Integer
n) [Atom]
xs

isPos :: Atom -> Bool
isPos :: Atom -> Bool
isPos Pos{} = Bool
True
isPos Neg{} = Bool
False

addAtom :: TTerm -> Atom -> TTerm
addAtom :: TTerm -> Atom -> TTerm
addAtom t :: TTerm
t (Pos a :: TTerm
a) = TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PAdd TTerm
t TTerm
a
addAtom t :: TTerm
t (Neg a :: TTerm
a) = TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PSub TTerm
t TTerm
a

toArith :: TTerm -> Arith
toArith :: TTerm -> Arith
toArith t :: TTerm
t | Just n :: Integer
n <- TTerm -> Maybe Integer
intView TTerm
t = (Integer
n, [])
toArith (TApp (TPrim PAdd) [a :: TTerm
a, b :: TTerm
b]) = Arith -> Arith -> Arith
aAdd (TTerm -> Arith
toArith TTerm
a) (TTerm -> Arith
toArith TTerm
b)
toArith (TApp (TPrim PSub) [a :: TTerm
a, b :: TTerm
b]) = Arith -> Arith -> Arith
aSub (TTerm -> Arith
toArith TTerm
a) (TTerm -> Arith
toArith TTerm
b)
toArith t :: TTerm
t = (0, [TTerm -> Atom
Pos TTerm
t])

simplArith :: TTerm -> TTerm
simplArith :: TTerm -> TTerm
simplArith = Arith -> TTerm
fromArith (Arith -> TTerm) -> (TTerm -> Arith) -> TTerm -> TTerm
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TTerm -> Arith
toArith