{-# OPTIONS_GHC -Wunused-imports #-}

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

import Control.Arrow        ( (***), second )
import Control.Monad        ( (>=>), guard )
import Control.Monad.Reader ( MonadReader(..), asks, Reader, runReader )
import qualified Data.List as List

import Agda.Syntax.Treeless
import Agda.Syntax.Literal

import Agda.TypeChecking.Monad
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 :: forall a. S a -> a
runS 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 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. EndoSubst 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 :: forall a.
(Substitution' TTerm -> Substitution' TTerm) -> S a -> S a
onSubst Substitution' TTerm -> Substitution' TTerm
f = (SEnv -> SEnv)
-> ReaderT SEnv Identity a -> ReaderT SEnv Identity a
forall a.
(SEnv -> SEnv)
-> ReaderT SEnv Identity a -> ReaderT SEnv Identity a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((SEnv -> SEnv)
 -> ReaderT SEnv Identity a -> ReaderT SEnv Identity a)
-> (SEnv -> SEnv)
-> ReaderT SEnv Identity a
-> ReaderT SEnv Identity a
forall a b. (a -> b) -> a -> b
$ \ SEnv
env -> SEnv
env { envSubst = f (envSubst env) }

onRewrite :: Substitution' TTerm -> S a -> S a
onRewrite :: forall a. Substitution' TTerm -> S a -> S a
onRewrite Substitution' TTerm
rho = (SEnv -> SEnv)
-> ReaderT SEnv Identity a -> ReaderT SEnv Identity a
forall a.
(SEnv -> SEnv)
-> ReaderT SEnv Identity a -> ReaderT SEnv Identity a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((SEnv -> SEnv)
 -> ReaderT SEnv Identity a -> ReaderT SEnv Identity a)
-> (SEnv -> SEnv)
-> ReaderT SEnv Identity a
-> ReaderT SEnv Identity a
forall a b. (a -> b) -> a -> b
$ \ SEnv
env -> SEnv
env { envRewrite = map (applySubst rho *** applySubst rho) (envRewrite env) }

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

underLams :: Int -> S a -> S a
underLams :: forall a. Int -> S a -> S a
underLams 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 :: forall a. S a -> S a
underLam = Int -> S a -> S a
forall a. Int -> S a -> S a
underLams Int
1

underLet :: TTerm -> S a -> S a
underLet :: forall a. TTerm -> S a -> S a
underLet 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 Int
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 (\Substitution' TTerm
rho -> Int -> Substitution' TTerm -> Substitution' TTerm
forall a. Int -> Substitution' a -> Substitution' a
wkS Int
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 :: forall a. Int -> TTerm -> S a -> S a
bindVar Int
x 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. EndoSubst a => Int -> a -> Substitution' a
inplaceS Int
x TTerm
u Substitution' TTerm -> Substitution' TTerm -> Substitution' TTerm
forall a.
EndoSubst a =>
Substitution' a -> Substitution' a -> Substitution' a
`composeS`)

rewrite :: TTerm -> S TTerm
rewrite :: TTerm -> S TTerm
rewrite TTerm
t = do
  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 [ rhs | (lhs, rhs) <- rules, equalTerms t lhs ] of
    TTerm
rhs : [TTerm]
_ -> TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
rhs
    []      -> TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
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 TTerm
t = do
  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
<$> BuiltinId -> TCMT IO (Maybe QName)
forall (m :: * -> *).
(HasBuiltins m, MonadReduce m) =>
BuiltinId -> m (Maybe QName)
getBuiltinName BuiltinId
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 a b. TCMT IO (a -> b) -> TCMT IO a -> TCMT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BuiltinId -> TCMT IO (Maybe QName)
forall (m :: * -> *).
(HasBuiltins m, MonadReduce m) =>
BuiltinId -> m (Maybe QName)
getBuiltinName BuiltinId
builtinNatDivSucAux
                     TCMT IO (Maybe QName -> Maybe QName -> Maybe QName -> FunctionKit)
-> TCMT IO (Maybe QName)
-> TCMT IO (Maybe QName -> Maybe QName -> FunctionKit)
forall a b. TCMT IO (a -> b) -> TCMT IO a -> TCMT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BuiltinId -> TCMT IO (Maybe QName)
forall (m :: * -> *).
(HasBuiltins m, MonadReduce m) =>
BuiltinId -> m (Maybe QName)
getBuiltinName BuiltinId
builtinNatMinus
                     TCMT IO (Maybe QName -> Maybe QName -> FunctionKit)
-> TCMT IO (Maybe QName) -> TCMT IO (Maybe QName -> FunctionKit)
forall a b. TCMT IO (a -> b) -> TCMT IO a -> TCMT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BuiltinId -> TCMT IO (Maybe QName)
forall (m :: * -> *).
(HasBuiltins m, MonadReduce m) =>
BuiltinId -> m (Maybe QName)
getBuiltinName BuiltinId
builtinTrue
                     TCMT IO (Maybe QName -> FunctionKit)
-> TCMT IO (Maybe QName) -> TCMT IO FunctionKit
forall a b. TCMT IO (a -> b) -> TCMT IO a -> TCMT IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BuiltinId -> TCMT IO (Maybe QName)
forall (m :: * -> *).
(HasBuiltins m, MonadReduce m) =>
BuiltinId -> m (Maybe QName)
getBuiltinName BuiltinId
builtinFalse
  return $ runS $ simplify kit t

simplify :: FunctionKit -> TTerm -> S TTerm
simplify :: FunctionKit -> TTerm -> S TTerm
simplify FunctionKit{Maybe QName
modAux :: FunctionKit -> Maybe QName
divAux :: FunctionKit -> Maybe QName
natMinus :: FunctionKit -> Maybe QName
true :: FunctionKit -> Maybe QName
false :: FunctionKit -> Maybe QName
modAux :: Maybe QName
divAux :: Maybe QName
natMinus :: Maybe QName
true :: Maybe QName
false :: Maybe QName
..} = 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 a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      t :: TTerm
t@TPrim{} -> TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      t :: TTerm
t@TVar{}  -> TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t

      TApp (TDef QName
f) [TLit (LitNat Integer
0), TTerm
m, TTerm
n, 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 Integer
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 Integer
1 TTerm
m)

      -- Word64 primitives --

      --  toWord (a ∙ b) == toWord a ∙64 toWord b
      TPFn TPrim
PITo64 (TPOp TPrim
op TTerm
a TTerm
b)
        | Just 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 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 TPrim
_) [TTerm]
_) -> TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t  -- taken care of by rewrite'

      TCoerce 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 TTerm
f [TTerm]
es -> do
        f  <- TTerm -> S TTerm
simpl TTerm
f
        es <- traverse simpl es
        maybeMinusToPrim f es
      TLam 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 a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      t :: TTerm
t@TCon{}  -> TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
      TLet TTerm
e TTerm
b  -> do
        TTerm -> S TTerm
simpl TTerm
e S TTerm -> (TTerm -> S TTerm) -> S TTerm
forall a b.
ReaderT SEnv Identity a
-> (a -> ReaderT SEnv Identity b) -> ReaderT SEnv Identity b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
          TPFn TPrim
P64ToI 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. EndoSubst a => Int -> a -> Substitution' a
inplaceS Int
0 (TPrim -> TTerm -> TTerm
TPFn TPrim
P64ToI (Int -> TTerm
TVar Int
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' (SubstArg TTerm) -> TTerm -> TTerm
forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' TTerm
Substitution' (SubstArg TTerm)
rho TTerm
b))
          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 Int
x CaseInfo
t TTerm
d [TAlt]
bs -> do
        v <- Int -> S TTerm
lookupVar Int
x
        let (lets, u) = tLetView v
        (d, bs) <- pruneBoolGuards d <$> traverse (simplAlt x) bs
        case u of                          -- TODO: also for literals
          TTerm
_ | Just (QName
c, [TTerm]
as)     <- TTerm -> Maybe (QName, [TTerm])
conView TTerm
u   -> TTerm -> S TTerm
simpl (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ [TTerm] -> QName -> [TTerm] -> TTerm -> [TAlt] -> TTerm
forall {t :: * -> *}.
Foldable t =>
t TTerm -> QName -> [TTerm] -> TTerm -> [TAlt] -> TTerm
matchCon [TTerm]
lets QName
c [TTerm]
as TTerm
d [TAlt]
bs
            | Just (Integer
k, TVar 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
. [TTerm] -> TTerm -> TTerm
forall {t :: * -> *}. Foldable t => t TTerm -> TTerm -> TTerm
mkLets [TTerm]
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)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM (Int -> Int -> Integer -> TAlt -> ReaderT SEnv Identity TAlt
matchPlusK Int
y Int
x Integer
k) [TAlt]
bs
          TCase Int
y CaseInfo
t1 TTerm
d1 [TAlt]
bs1 -> TTerm -> S TTerm
simpl (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ [TTerm] -> TTerm -> TTerm
forall {t :: * -> *}. Foldable t => t TTerm -> TTerm -> TTerm
mkLets [TTerm]
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     = [TTerm] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TTerm]
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
+ Int
1) (Int -> Substitution' TTerm
forall a. Int -> Substitution' a
raiseS Int
1)    Substitution' TTerm -> Substitution' TTerm -> Substitution' TTerm
forall a.
EndoSubst 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
+ Int
1) (Int -> TTerm
TVar Int
0) Substitution' TTerm -> Substitution' TTerm -> Substitution' TTerm
forall a.
EndoSubst 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
+ Int
1)
              case1 :: TTerm
case1 = Substitution' (SubstArg TTerm) -> TTerm -> TTerm
forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' TTerm
Substitution' (SubstArg TTerm)
rho (Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
x CaseInfo
t TTerm
d [TAlt]
bs)

              distrDef :: TTerm -> TTerm -> TTerm
distrDef TTerm
v 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 TTerm
v (TACon QName
c Int
a 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 a. Subst a => Int -> Int -> a -> a
raiseFrom Int
1 Int
a TTerm
v
              distrCase TTerm
v (TALit Literal
l 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 TTerm
v (TAGuard TTerm
g 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

          TTerm
_ -> do
            d <- TTerm -> S TTerm
simpl TTerm
d
            tCase x t d bs

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

    conView :: TTerm -> Maybe (QName, [TTerm])
conView (TCon QName
c)    = (QName, [TTerm]) -> Maybe (QName, [TTerm])
forall a. a -> Maybe a
Just (QName
c, [])
    conView (TApp TTerm
f [TTerm]
as) = ([TTerm] -> [TTerm]) -> (QName, [TTerm]) -> (QName, [TTerm])
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second ([TTerm] -> [TTerm] -> [TTerm]
forall a. [a] -> [a] -> [a]
++ [TTerm]
as) ((QName, [TTerm]) -> (QName, [TTerm]))
-> Maybe (QName, [TTerm]) -> Maybe (QName, [TTerm])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> Maybe (QName, [TTerm])
conView TTerm
f
    conView TTerm
e           = Maybe (QName, [TTerm])
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 Int
x CaseInfo
t TTerm
d [TAlt]
bs) = do
      let ([TTerm]
lets, TTerm
u) = TTerm -> ([TTerm], TTerm)
tLetView TTerm
d
          k :: Int
k = [TTerm] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [TTerm]
lets
      TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
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 Int
y CaseInfo
_ TTerm
d' [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 ->
          [TTerm] -> TTerm -> TTerm
forall {t :: * -> *}. Foldable t => t TTerm -> TTerm -> TTerm
mkLets [TTerm]
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 a. Subst a => Int -> a -> a
raise Int
k [TAlt]
bs [TAlt] -> [TAlt] -> [TAlt]
forall a. [a] -> [a] -> [a]
++ [TAlt]
bs'
        TTerm
_ -> TTerm
e
    unchainCase TTerm
e = TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
e


    mkLets :: t TTerm -> TTerm -> TTerm
mkLets t TTerm
es TTerm
b = (TTerm -> TTerm -> TTerm) -> TTerm -> t TTerm -> TTerm
forall a b. (a -> b -> b) -> b -> t a -> b
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 -> [TTerm] -> TTerm -> [TAlt] -> TTerm
matchCon t TTerm
_ QName
_ [TTerm]
_ TTerm
d [] = TTerm
d
    matchCon t TTerm
lets QName
c [TTerm]
as TTerm
d (TALit{}   : [TAlt]
bs) = t TTerm -> QName -> [TTerm] -> TTerm -> [TAlt] -> TTerm
matchCon t TTerm
lets QName
c [TTerm]
as TTerm
d [TAlt]
bs
    matchCon t TTerm
lets QName
c [TTerm]
as TTerm
d (TAGuard{} : [TAlt]
bs) = t TTerm -> QName -> [TTerm] -> TTerm -> [TAlt] -> TTerm
matchCon t TTerm
lets QName
c [TTerm]
as TTerm
d [TAlt]
bs
    matchCon t TTerm
lets QName
c [TTerm]
as TTerm
d (TACon QName
c' Int
a TTerm
b : [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 a b. (a -> b -> b) -> b -> t a -> b
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 -> [TTerm] -> TTerm -> TTerm
mkLet Int
0 [TTerm]
as (Int -> Int -> TTerm -> TTerm
forall a. Subst a => Int -> Int -> a -> a
raiseFrom Int
a (t TTerm -> Int
forall a. t a -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length t TTerm
lets) TTerm
b)
      | Bool
otherwise      = t TTerm -> QName -> [TTerm] -> TTerm -> [TAlt] -> TTerm
matchCon t TTerm
lets QName
c [TTerm]
as TTerm
d [TAlt]
bs
      where
        mkLet :: Int -> [TTerm] -> TTerm -> TTerm
mkLet Int
_ []       TTerm
b = TTerm
b
        mkLet Int
i (TTerm
a : [TTerm]
as) TTerm
b = TTerm -> TTerm -> TTerm
TLet (Int -> TTerm -> TTerm
forall a. Subst a => Int -> a -> a
raise Int
i TTerm
a) (TTerm -> TTerm) -> TTerm -> TTerm
forall a b. (a -> b) -> a -> b
$ Int -> [TTerm] -> TTerm -> TTerm
mkLet (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [TTerm]
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 Int
x Int
y Integer
k (TALit (LitNat Integer
j) TTerm
b) = TAlt -> ReaderT SEnv Identity TAlt
forall a. a -> ReaderT SEnv Identity a
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 (Integer -> Literal
LitNat (Integer
j Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
k)) TTerm
b
    matchPlusK Int
x Int
y Integer
k (TAGuard TTerm
g 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' (SubstArg TTerm) -> TTerm -> TTerm
forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst (Int -> TTerm -> Substitution' TTerm
forall a. EndoSubst a => Int -> a -> Substitution' a
inplaceS Int
y (Integer -> TTerm -> TTerm
tPlusK Integer
k (Int -> TTerm
TVar Int
x))) TTerm
g)
    matchPlusK Int
x Int
y Integer
k TACon{} = ReaderT SEnv Identity TAlt
forall a. HasCallStack => a
__IMPOSSIBLE__
    matchPlusK Int
x Int
y Integer
k TALit{} = ReaderT SEnv Identity TAlt
forall a. HasCallStack => a
__IMPOSSIBLE__

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

    simplPrim' :: TTerm -> TTerm
    simplPrim' :: TTerm -> TTerm
simplPrim' (TApp (TPrim TPrim
PSeq) (TTerm
u : TTerm
v : [TTerm]
vs))
      | TTerm
u TTerm -> TTerm -> Bool
forall a. Eq a => a -> a -> Bool
== TTerm
v             = TTerm -> [TTerm] -> TTerm
mkTApp TTerm
v [TTerm]
vs
      | TApp TCon{} [TTerm]
_ <- TTerm
u = TTerm -> [TTerm] -> TTerm
mkTApp TTerm
v [TTerm]
vs
      | TApp TLit{} [TTerm]
_ <- TTerm
u = TTerm -> [TTerm] -> TTerm
mkTApp TTerm
v [TTerm]
vs
    simplPrim' (TApp (TPrim TPrim
PLt) [TTerm
u, TTerm
v])
      | Just (TPrim
PAdd, Integer
k, TTerm
u) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
u,
        Just (TPrim
PAdd, Integer
j, 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 (TPrim
PSub, Integer
k, TTerm
u) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
u,
        Just (TPrim
PSub, Integer
j, 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
v TTerm
u
      | Just (TPrim
PAdd, Integer
k, TTerm
v) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
v,
        TApp (TPrim TPrim
P64ToI) [TTerm
u] <- TTerm
u,
        Integer
k Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer
2 Integer -> Integer -> Integer
forall a b. (Num a, Integral b) => a -> b -> a
^ Integer
64, Just QName
trueCon <- Maybe QName
true = QName -> TTerm
TCon QName
trueCon
      | Just Integer
k <- TTerm -> Maybe Integer
intView TTerm
u
      , Just Integer
j <- TTerm -> Maybe Integer
intView TTerm
v
      , Just QName
trueCon <- Maybe QName
true
      , Just 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 TPrim
PGeq) [TTerm
u, TTerm
v])
      | Just (TPrim
PAdd, Integer
k, TTerm
u) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
u,
        Just (TPrim
PAdd, Integer
j, 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
PGeq TTerm
u TTerm
v
      | Just (TPrim
PSub, Integer
k, TTerm
u) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
u,
        Just (TPrim
PSub, Integer
j, 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
PGeq TTerm
v TTerm
u
      | Just Integer
k <- TTerm -> Maybe Integer
intView TTerm
u
      , Just Integer
j <- TTerm -> Maybe Integer
intView TTerm
v
      , Just QName
trueCon <- Maybe QName
true
      , Just 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 TPrim
op) [TTerm
u, TTerm
v])
      | TPrim
op TPrim -> [TPrim] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PGeq, TPrim
PLt, TPrim
PEqI]
      , Just (TPrim
PAdd, Integer
k, TTerm
u) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
u
      , Just Integer
j <- TTerm -> Maybe Integer
intView TTerm
v = TTerm -> [TTerm] -> 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 TPrim
PEqI) [TTerm
u, TTerm
v])
      | Just (TPrim
op1, Integer
k, TTerm
u) <- TTerm -> Maybe (TPrim, Integer, TTerm)
constArithView TTerm
u,
        Just (TPrim
op2, Integer
j, 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 a. Eq a => a -> [a] -> 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 TPrim
op TTerm
u TTerm
v)
      | Bool
zeroL, Bool
isMul Bool -> Bool -> Bool
|| Bool
isDiv = Integer -> TTerm
tInt Integer
0
      | Bool
zeroL, Bool
isAdd          = TTerm
v
      | Bool
zeroR, Bool
isMul          = Integer -> TTerm
tInt Integer
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 Integer
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 Word64
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 Integer
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 Word64
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 a. Eq a => a -> [a] -> 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 a. Eq a => a -> [a] -> 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 a. Eq a => a -> [a] -> 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 a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [TPrim
PQuot, TPrim
PQuot64, TPrim
PRem, TPrim
PRem64]
    simplPrim' (TApp (TPrim TPrim
op) [TTerm
u, TTerm
v])
      | Just TTerm
u <- TTerm -> Maybe TTerm
negView TTerm
u,
        Just TTerm
v <- TTerm -> Maybe TTerm
negView TTerm
v,
        TPrim
op TPrim -> [TPrim] -> Bool
forall a. Eq a => a -> [a] -> 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 TTerm
u <- TTerm -> Maybe TTerm
negView TTerm
u,
        TPrim
op TPrim -> [TPrim] -> Bool
forall a. Eq a => a -> [a] -> 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 Integer
0) (TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
op TTerm
u TTerm
v)
      | Just TTerm
v <- TTerm -> Maybe TTerm
negView TTerm
v,
        TPrim
op TPrim -> [TPrim] -> Bool
forall a. Eq a => a -> [a] -> 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 Integer
0) (TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
op TTerm
u TTerm
v)
    simplPrim' (TApp (TPrim TPrim
PRem) [TTerm
u, TTerm
v])
      | Just 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 Integer
0) (TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PRem TTerm
u (TTerm -> TTerm
unNeg TTerm
v))
      | Just 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 TPrim
op (TPFn TPrim
P64ToI TTerm
a) (TPFn TPrim
P64ToI TTerm
b))
        | Just 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 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 TPrim
PITo64 (TLit (LitNat Integer
n)))    = Literal -> TTerm
TLit (Word64 -> Literal
LitWord64 (Integer -> Word64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n))
    simplPrim' (TPFn TPrim
P64ToI (TLit (LitWord64 Word64
n))) = Literal -> TTerm
TLit (Integer -> Literal
LitNat    (Word64 -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
n))

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

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

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

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

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

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

    rewrite' :: TTerm -> S TTerm
rewrite' 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 TPrim
op) [TLit (LitNat Integer
k), TTerm
u])
      | TPrim
op TPrim -> [TPrim] -> Bool
forall a. Eq a => a -> [a] -> 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 TPrim
op) [TTerm
u, TLit (LitNat 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 TTerm
_ = Maybe (TPrim, Integer, TTerm)
forall a. Maybe a
Nothing

    simplAlt :: Int -> TAlt -> ReaderT SEnv Identity TAlt
simplAlt Int
x (TACon QName
c Int
a 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 -> [TTerm] -> TTerm
mkTApp (QName -> TTerm
TCon QName
c) ([TTerm] -> TTerm) -> [TTerm] -> TTerm
forall a b. (a -> b) -> a -> b
$ (Int -> TTerm) -> [Int] -> [TTerm]
forall a b. (a -> b) -> [a] -> [b]
map Int -> TTerm
TVar ([Int] -> [TTerm]) -> [Int] -> [TTerm]
forall a b. (a -> b) -> a -> b
$ Int -> [Int]
forall a. Integral a => a -> [a]
downFrom Int
a
    simplAlt Int
x (TALit Literal
l 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 Int
x (TAGuard TTerm
g 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 a b.
ReaderT SEnv Identity (a -> b)
-> ReaderT SEnv Identity a -> ReaderT SEnv Identity b
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 b -> S b
maybeAddRewrite Int
x TTerm
rhs S b
cont = do
      v <- Int -> S TTerm
lookupVar Int
x
      case v of
        TVar Int
y | Int
x Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
y -> Int -> TTerm -> S b -> S b
forall a. Int -> TTerm -> S a -> S a
bindVar Int
x TTerm
rhs (S b -> S b) -> S b -> S b
forall a b. (a -> b) -> a -> b
$ S b
cont
        TTerm
_ -> TTerm -> TTerm -> S b -> S b
forall a. TTerm -> TTerm -> S a -> S a
addRewrite TTerm
v TTerm
rhs S b
cont

    isTrue :: TTerm -> Bool
isTrue (TCon 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 TTerm
_        = Bool
False

    isFalse :: TTerm -> Bool
isFalse (TCon 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 TTerm
_        = Bool
False

    maybeMinusToPrim :: TTerm -> [TTerm] -> S TTerm
maybeMinusToPrim f :: TTerm
f@(TDef QName
minus) es :: [TTerm]
es@[TTerm
a, 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
      leq  <- TTerm -> TTerm -> ReaderT SEnv Identity Bool
checkLeq TTerm
b TTerm
a
      if leq then pure $ tOp PSub a b
             else tApp f es

    maybeMinusToPrim TTerm
f [TTerm]
es = TTerm -> [TTerm] -> S TTerm
tApp TTerm
f [TTerm]
es

    tLet :: TTerm -> TTerm -> TTerm
tLet (TVar Int
x) TTerm
b = Int -> SubstArg TTerm -> TTerm -> TTerm
forall a. Subst a => Int -> SubstArg a -> a -> a
subst Int
0 (Int -> TTerm
TVar Int
x) TTerm
b
    tLet TTerm
e (TVar Int
0) = TTerm
e
    tLet TTerm
e 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 Int
x CaseInfo
t TTerm
d [] = TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
d
    tCase Int
x CaseInfo
t TTerm
d [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 a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
d
          TALit Literal
_ TTerm
b   : [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 TTerm
_ TTerm
b : [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 QName
c Int
a TTerm
b : [TAlt]
_   -> Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase' Int
x CaseInfo
t TTerm
d [TAlt]
bs'
      | Bool
otherwise = do
        d' <- TTerm -> S TTerm
lookupIfVar TTerm
d
        case d' of
          TCase Int
y CaseInfo
_ TTerm
d [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'')
          TTerm
_ -> 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 Int
i) = Int -> S TTerm
lookupVar Int
i
        lookupIfVar TTerm
t = TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t

        noOverlap :: TAlt -> Bool
noOverlap 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 QName
c Int
_ TTerm
_)  (TACon QName
c' Int
_ TTerm
_) = QName
c QName -> QName -> Bool
forall a. Eq a => a -> a -> Bool
== QName
c'
        overlapped (TALit Literal
l TTerm
_)    (TALit Literal
l' TTerm
_)   = Literal
l Literal -> Literal -> Bool
forall a. Eq a => a -> a -> Bool
== Literal
l'
        overlapped TAlt
_              TAlt
_              = 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 Int
x CaseInfo
t TTerm
d [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 [TAlt]
bs' -> Int -> CaseInfo -> TTerm -> [TAlt] -> S TTerm
tCase Int
x CaseInfo
t TTerm
tUnreachable [TAlt]
bs'
        Maybe [TAlt]
Nothing  -> TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
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 [TAlt]
bs [Integer]
small (Just Integer
upper)
          | [Integer] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null ([Integer] -> Bool) -> [Integer] -> Bool
forall a b. (a -> b) -> a -> b
$ [Integer
0..Integer
upper Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
- Integer
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 Integer
n) TTerm
_) : [TAlt]
bs) [Integer]
small 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 TPrim
PGeq) [TVar Int
y, TLit (LitNat Integer
j)]) TTerm
_) : [TAlt]
bs) [Integer]
small 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 [TAlt]
_ [Integer]
_ Maybe Integer
_ = Maybe [TAlt]
forall a. Maybe a
Nothing

    pruneLitCases Int
x CaseInfo
t TTerm
d [TAlt]
bs
      | CaseType
CTInt CaseType -> CaseType -> Bool
forall a. Eq a => a -> a -> Bool
== CaseInfo -> CaseType
caseType CaseInfo
t = TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
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 a. a -> ReaderT SEnv Identity a
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

    -- Drop 'false' branches and drop everything after 'true' branches (including the default
    -- branch)
    pruneBoolGuards :: TTerm -> [TAlt] -> (TTerm, [TAlt])
pruneBoolGuards TTerm
d [] = (TTerm
d, [])
    pruneBoolGuards TTerm
d (b :: TAlt
b@(TAGuard (TCon QName
c) TTerm
_) : [TAlt]
bs)
      | 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  = (TTerm
tUnreachable, [TAlt
b])
      | 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 = TTerm -> [TAlt] -> (TTerm, [TAlt])
pruneBoolGuards TTerm
d [TAlt]
bs
    pruneBoolGuards TTerm
d (TAlt
b : [TAlt]
bs) =
      ([TAlt] -> [TAlt]) -> (TTerm, [TAlt]) -> (TTerm, [TAlt])
forall b c d. (b -> c) -> (d, b) -> (d, c)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (d, b) (d, c)
second (TAlt
b TAlt -> [TAlt] -> [TAlt]
forall a. a -> [a] -> [a]
:) ((TTerm, [TAlt]) -> (TTerm, [TAlt]))
-> (TTerm, [TAlt]) -> (TTerm, [TAlt])
forall a b. (a -> b) -> a -> b
$ TTerm -> [TAlt] -> (TTerm, [TAlt])
pruneBoolGuards TTerm
d [TAlt]
bs

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

    tApp :: TTerm -> [TTerm] -> S TTerm
    tApp :: TTerm -> [TTerm] -> S TTerm
tApp (TLet TTerm
e TTerm
b) [TTerm]
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 -> [TTerm] -> S TTerm
tApp TTerm
b (Int -> [TTerm] -> [TTerm]
forall a. Subst a => Int -> a -> a
raise Int
1 [TTerm]
es))
    tApp (TCase Int
x CaseInfo
t TTerm
d [TAlt]
bs) [TTerm]
es = do
      d  <- TTerm -> [TTerm] -> S TTerm
tApp TTerm
d [TTerm]
es
      bs <- mapM (`tAppAlt` es) bs
      simpl $ TCase x t d bs    -- will resimplify branches
    tApp (TVar Int
x) [TTerm]
es = do
      v <- Int -> S TTerm
lookupVar Int
x
      case v of
        TTerm
_ | 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 -> [TTerm] -> S TTerm
tApp TTerm
v [TTerm]
es
        TLam{} -> TTerm -> [TTerm] -> S TTerm
tApp TTerm
v [TTerm]
es   -- could blow up the code
        TTerm
_      -> TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> [TTerm] -> TTerm
mkTApp (Int -> TTerm
TVar Int
x) [TTerm]
es
    tApp TTerm
f [] = TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
f
    tApp (TLam TTerm
b) (TVar Int
i : [TTerm]
es) = TTerm -> [TTerm] -> S TTerm
tApp (Int -> SubstArg TTerm -> TTerm -> TTerm
forall a. Subst a => Int -> SubstArg a -> a -> a
subst Int
0 (Int -> TTerm
TVar Int
i) TTerm
b) [TTerm]
es
    tApp (TLam TTerm
b) (TTerm
e : [TTerm]
es) = TTerm -> [TTerm] -> S TTerm
tApp (TTerm -> TTerm -> TTerm
TLet TTerm
e TTerm
b) [TTerm]
es
    tApp TTerm
f [TTerm]
es = TTerm -> S TTerm
forall a. a -> ReaderT SEnv Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> [TTerm] -> TTerm
TApp TTerm
f [TTerm]
es

    tAppAlt :: TAlt -> [TTerm] -> ReaderT SEnv Identity TAlt
tAppAlt (TACon QName
c Int
a TTerm
b) [TTerm]
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 -> [TTerm] -> S TTerm
tApp TTerm
b (Int -> [TTerm] -> [TTerm]
forall a. Subst a => Int -> a -> a
raise Int
a [TTerm]
es))
    tAppAlt (TALit Literal
l TTerm
b) [TTerm]
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 -> [TTerm] -> S TTerm
tApp TTerm
b [TTerm]
es
    tAppAlt (TAGuard TTerm
g TTerm
b) [TTerm]
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 -> [TTerm] -> S TTerm
tApp TTerm
b [TTerm]
es

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

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

          match (a
j, a
as) (a
k, 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
d (a, a)
x (a, a)
y ((a, a)
x', (a, a)
y') = Maybe () -> Bool
forall a. Maybe a -> Bool
isJust (Maybe () -> Bool) -> Maybe () -> Bool
forall a b. (a -> b) -> a -> b
$ do
            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
            j <- match y y'     -- y = y' + j
            guard (k <= j + d)  -- x ≤ y if k ≤ j + d

          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 Integer
1
          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 Integer
0

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

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

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

          go Arith
x 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 ]

      return $ go (nf a) (nf 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
$cshowsPrec :: Int -> Atom -> ShowS
showsPrec :: Int -> Atom -> ShowS
$cshow :: Atom -> String
show :: Atom -> String
$cshowList :: [Atom] -> ShowS
showList :: [Atom] -> ShowS
Show, Atom -> Atom -> Bool
(Atom -> Atom -> Bool) -> (Atom -> Atom -> Bool) -> Eq Atom
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Atom -> Atom -> Bool
== :: Atom -> Atom -> Bool
$c/= :: Atom -> Atom -> Bool
/= :: 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
$ccompare :: Atom -> Atom -> Ordering
compare :: Atom -> Atom -> Ordering
$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
>= :: Atom -> Atom -> Bool
$cmax :: Atom -> Atom -> Atom
max :: Atom -> Atom -> Atom
$cmin :: Atom -> Atom -> Atom
min :: Atom -> Atom -> Atom
Ord)

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

aCancel :: [Atom] -> [Atom]
aCancel :: [Atom] -> [Atom]
aCancel (Atom
a : [Atom]
as)
  | (Atom -> Atom
aNeg Atom
a) Atom -> [Atom] -> Bool
forall a. Eq a => a -> [a] -> 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 :: forall a. Ord a => [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 (Integer
a, [Atom]
xs) (Integer
b, [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 (Integer
a, [Atom]
xs) (Integer
b, [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 (Integer
n, []) = Integer -> TTerm
tInt Integer
n
fromArith (Integer
0, [Atom]
xs)
  | ([Atom]
ys, Pos TTerm
a : [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 b a. (b -> a -> b) -> b -> [a] -> b
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 (Integer
n, [Atom]
xs)
  | Integer
n Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Integer
0, ([Atom]
ys, Pos TTerm
a : [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 b a. (b -> a -> b) -> b -> [a] -> b
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 (Integer
n, [Atom]
xs) = (TTerm -> Atom -> TTerm) -> TTerm -> [Atom] -> TTerm
forall b a. (b -> a -> b) -> b -> [a] -> b
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 TTerm
t (Pos TTerm
a) = TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PAdd TTerm
t TTerm
a
addAtom TTerm
t (Neg TTerm
a) = TPrim -> TTerm -> TTerm -> TTerm
tOp TPrim
PSub TTerm
t TTerm
a

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