{-# OPTIONS_GHC -fno-warn-orphans       #-}

module Agda.Compiler.Treeless.Subst where

import qualified Data.Map as Map
import Data.Map (Map)
import Data.Maybe
import Data.Semigroup ( Semigroup, (<>), All(..), Any(..) )

import Agda.Syntax.Treeless
import Agda.TypeChecking.Substitute

import Agda.Utils.Impossible

instance DeBruijn TTerm where
  deBruijnVar :: Int -> TTerm
deBruijnVar = Int -> TTerm
TVar
  deBruijnView :: TTerm -> Maybe Int
deBruijnView (TVar Int
i) = forall a. a -> Maybe a
Just Int
i
  deBruijnView TTerm
_ = forall a. Maybe a
Nothing

instance Subst TTerm where
  type SubstArg TTerm = TTerm

  applySubst :: Substitution' (SubstArg TTerm) -> TTerm -> TTerm
applySubst Substitution' (SubstArg TTerm)
IdS = forall a. a -> a
id
  applySubst Substitution' (SubstArg TTerm)
rho = \case
      t :: TTerm
t@TDef{}    -> TTerm
t
      t :: TTerm
t@TLit{}    -> TTerm
t
      t :: TTerm
t@TCon{}    -> TTerm
t
      t :: TTerm
t@TPrim{}   -> TTerm
t
      t :: TTerm
t@TUnit{}   -> TTerm
t
      t :: TTerm
t@TSort{}   -> TTerm
t
      t :: TTerm
t@TErased{} -> TTerm
t
      t :: TTerm
t@TError{}  -> TTerm
t
      TVar Int
i         -> forall a. EndoSubst a => Substitution' a -> Int -> a
lookupS Substitution' (SubstArg TTerm)
rho Int
i
      TApp TTerm
f [TTerm]
ts      -> TTerm -> [TTerm] -> TTerm
tApp (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TTerm)
rho TTerm
f) (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TTerm)
rho [TTerm]
ts)
      TLam TTerm
b         -> TTerm -> TTerm
TLam (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst (forall a. Int -> Substitution' a -> Substitution' a
liftS Int
1 Substitution' (SubstArg TTerm)
rho) TTerm
b)
      TLet TTerm
e TTerm
b       -> TTerm -> TTerm -> TTerm
TLet (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TTerm)
rho TTerm
e) (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst (forall a. Int -> Substitution' a -> Substitution' a
liftS Int
1 Substitution' (SubstArg TTerm)
rho) TTerm
b)
      TCase Int
i CaseInfo
t TTerm
d [TAlt]
bs ->
        case forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TTerm)
rho (Int -> TTerm
TVar Int
i) of
          TVar Int
j  -> Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
j CaseInfo
t (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TTerm)
rho TTerm
d) (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TTerm)
rho [TAlt]
bs)
          TTerm
e       -> TTerm -> TTerm -> TTerm
TLet TTerm
e forall a b. (a -> b) -> a -> b
$ Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
0 CaseInfo
t (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' TTerm
rho' TTerm
d) (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' TTerm
rho' [TAlt]
bs)
            where rho' :: Substitution' TTerm
rho' = forall a. Int -> Substitution' a -> Substitution' a
wkS Int
1 Substitution' (SubstArg TTerm)
rho
      TCoerce TTerm
e -> TTerm -> TTerm
TCoerce (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TTerm)
rho TTerm
e)
    where
      tApp :: TTerm -> [TTerm] -> TTerm
tApp (TPrim TPrim
PSeq) [TTerm
TErased, TTerm
b] = TTerm
b
      tApp TTerm
f [TTerm]
ts = TTerm -> [TTerm] -> TTerm
TApp TTerm
f [TTerm]
ts

instance Subst TAlt where
  type SubstArg TAlt = TTerm
  applySubst :: Substitution' (SubstArg TAlt) -> TAlt -> TAlt
applySubst Substitution' (SubstArg TAlt)
rho (TACon QName
c Int
i TTerm
b) = QName -> Int -> TTerm -> TAlt
TACon QName
c Int
i (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst (forall a. Int -> Substitution' a -> Substitution' a
liftS Int
i Substitution' (SubstArg TAlt)
rho) TTerm
b)
  applySubst Substitution' (SubstArg TAlt)
rho (TALit Literal
l TTerm
b)   = Literal -> TTerm -> TAlt
TALit Literal
l (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TAlt)
rho TTerm
b)
  applySubst Substitution' (SubstArg TAlt)
rho (TAGuard TTerm
g TTerm
b) = TTerm -> TTerm -> TAlt
TAGuard (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TAlt)
rho TTerm
g) (forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst Substitution' (SubstArg TAlt)
rho TTerm
b)

newtype UnderLambda = UnderLambda Any
  deriving (UnderLambda -> UnderLambda -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: UnderLambda -> UnderLambda -> Bool
$c/= :: UnderLambda -> UnderLambda -> Bool
== :: UnderLambda -> UnderLambda -> Bool
$c== :: UnderLambda -> UnderLambda -> Bool
Eq, Eq UnderLambda
UnderLambda -> UnderLambda -> Bool
UnderLambda -> UnderLambda -> Ordering
UnderLambda -> UnderLambda -> UnderLambda
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 :: UnderLambda -> UnderLambda -> UnderLambda
$cmin :: UnderLambda -> UnderLambda -> UnderLambda
max :: UnderLambda -> UnderLambda -> UnderLambda
$cmax :: UnderLambda -> UnderLambda -> UnderLambda
>= :: UnderLambda -> UnderLambda -> Bool
$c>= :: UnderLambda -> UnderLambda -> Bool
> :: UnderLambda -> UnderLambda -> Bool
$c> :: UnderLambda -> UnderLambda -> Bool
<= :: UnderLambda -> UnderLambda -> Bool
$c<= :: UnderLambda -> UnderLambda -> Bool
< :: UnderLambda -> UnderLambda -> Bool
$c< :: UnderLambda -> UnderLambda -> Bool
compare :: UnderLambda -> UnderLambda -> Ordering
$ccompare :: UnderLambda -> UnderLambda -> Ordering
Ord, Int -> UnderLambda -> ShowS
[UnderLambda] -> ShowS
UnderLambda -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [UnderLambda] -> ShowS
$cshowList :: [UnderLambda] -> ShowS
show :: UnderLambda -> String
$cshow :: UnderLambda -> String
showsPrec :: Int -> UnderLambda -> ShowS
$cshowsPrec :: Int -> UnderLambda -> ShowS
Show, NonEmpty UnderLambda -> UnderLambda
UnderLambda -> UnderLambda -> UnderLambda
forall b. Integral b => b -> UnderLambda -> UnderLambda
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: forall b. Integral b => b -> UnderLambda -> UnderLambda
$cstimes :: forall b. Integral b => b -> UnderLambda -> UnderLambda
sconcat :: NonEmpty UnderLambda -> UnderLambda
$csconcat :: NonEmpty UnderLambda -> UnderLambda
<> :: UnderLambda -> UnderLambda -> UnderLambda
$c<> :: UnderLambda -> UnderLambda -> UnderLambda
Semigroup, Semigroup UnderLambda
UnderLambda
[UnderLambda] -> UnderLambda
UnderLambda -> UnderLambda -> UnderLambda
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [UnderLambda] -> UnderLambda
$cmconcat :: [UnderLambda] -> UnderLambda
mappend :: UnderLambda -> UnderLambda -> UnderLambda
$cmappend :: UnderLambda -> UnderLambda -> UnderLambda
mempty :: UnderLambda
$cmempty :: UnderLambda
Monoid)

newtype SeqArg = SeqArg All
  deriving (SeqArg -> SeqArg -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: SeqArg -> SeqArg -> Bool
$c/= :: SeqArg -> SeqArg -> Bool
== :: SeqArg -> SeqArg -> Bool
$c== :: SeqArg -> SeqArg -> Bool
Eq, Eq SeqArg
SeqArg -> SeqArg -> Bool
SeqArg -> SeqArg -> Ordering
SeqArg -> SeqArg -> SeqArg
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 :: SeqArg -> SeqArg -> SeqArg
$cmin :: SeqArg -> SeqArg -> SeqArg
max :: SeqArg -> SeqArg -> SeqArg
$cmax :: SeqArg -> SeqArg -> SeqArg
>= :: SeqArg -> SeqArg -> Bool
$c>= :: SeqArg -> SeqArg -> Bool
> :: SeqArg -> SeqArg -> Bool
$c> :: SeqArg -> SeqArg -> Bool
<= :: SeqArg -> SeqArg -> Bool
$c<= :: SeqArg -> SeqArg -> Bool
< :: SeqArg -> SeqArg -> Bool
$c< :: SeqArg -> SeqArg -> Bool
compare :: SeqArg -> SeqArg -> Ordering
$ccompare :: SeqArg -> SeqArg -> Ordering
Ord, Int -> SeqArg -> ShowS
[SeqArg] -> ShowS
SeqArg -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [SeqArg] -> ShowS
$cshowList :: [SeqArg] -> ShowS
show :: SeqArg -> String
$cshow :: SeqArg -> String
showsPrec :: Int -> SeqArg -> ShowS
$cshowsPrec :: Int -> SeqArg -> ShowS
Show, NonEmpty SeqArg -> SeqArg
SeqArg -> SeqArg -> SeqArg
forall b. Integral b => b -> SeqArg -> SeqArg
forall a.
(a -> a -> a)
-> (NonEmpty a -> a)
-> (forall b. Integral b => b -> a -> a)
-> Semigroup a
stimes :: forall b. Integral b => b -> SeqArg -> SeqArg
$cstimes :: forall b. Integral b => b -> SeqArg -> SeqArg
sconcat :: NonEmpty SeqArg -> SeqArg
$csconcat :: NonEmpty SeqArg -> SeqArg
<> :: SeqArg -> SeqArg -> SeqArg
$c<> :: SeqArg -> SeqArg -> SeqArg
Semigroup, Semigroup SeqArg
SeqArg
[SeqArg] -> SeqArg
SeqArg -> SeqArg -> SeqArg
forall a.
Semigroup a -> a -> (a -> a -> a) -> ([a] -> a) -> Monoid a
mconcat :: [SeqArg] -> SeqArg
$cmconcat :: [SeqArg] -> SeqArg
mappend :: SeqArg -> SeqArg -> SeqArg
$cmappend :: SeqArg -> SeqArg -> SeqArg
mempty :: SeqArg
$cmempty :: SeqArg
Monoid)

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

once :: Occurs
once :: Occurs
once = Int -> UnderLambda -> SeqArg -> Occurs
Occurs Int
1 forall a. Monoid a => a
mempty (All -> SeqArg
SeqArg forall a b. (a -> b) -> a -> b
$ Bool -> All
All Bool
False)

inSeq :: Occurs -> Occurs
inSeq :: Occurs -> Occurs
inSeq (Occurs Int
n UnderLambda
l SeqArg
_) = Int -> UnderLambda -> SeqArg -> Occurs
Occurs Int
n UnderLambda
l forall a. Monoid a => a
mempty

underLambda :: Occurs -> Occurs
underLambda :: Occurs -> Occurs
underLambda Occurs
o = Occurs
o forall a. Semigroup a => a -> a -> a
<> Int -> UnderLambda -> SeqArg -> Occurs
Occurs Int
0 (Any -> UnderLambda
UnderLambda forall a b. (a -> b) -> a -> b
$ Bool -> Any
Any Bool
True) forall a. Monoid a => a
mempty

instance Semigroup Occurs where
  Occurs Int
a UnderLambda
k SeqArg
s <> :: Occurs -> Occurs -> Occurs
<> Occurs Int
b UnderLambda
l SeqArg
t = Int -> UnderLambda -> SeqArg -> Occurs
Occurs (Int
a forall a. Num a => a -> a -> a
+ Int
b) (UnderLambda
k forall a. Semigroup a => a -> a -> a
<> UnderLambda
l) (SeqArg
s forall a. Semigroup a => a -> a -> a
<> SeqArg
t)

instance Monoid Occurs where
  mempty :: Occurs
mempty  = Int -> UnderLambda -> SeqArg -> Occurs
Occurs Int
0 forall a. Monoid a => a
mempty forall a. Monoid a => a
mempty
  mappend :: Occurs -> Occurs -> Occurs
mappend = forall a. Semigroup a => a -> a -> a
(<>)


-- Andreas, 2019-07-10: this free variable computation should be rewritten
-- in the style of TypeChecking.Free.Lazy.
-- https://github.com/agda/agda/commit/03eb3945114a4ccdb449f22d69db8d6eaa4699b8#commitcomment-34249120

class HasFree a where
  freeVars :: a -> Map Int Occurs

freeIn :: HasFree a => Int -> a -> Bool
freeIn :: forall a. HasFree a => Int -> a -> Bool
freeIn Int
i a
x = forall k a. Ord k => k -> Map k a -> Bool
Map.member Int
i (forall a. HasFree a => a -> Map Int Occurs
freeVars a
x)

occursIn :: HasFree a => Int -> a -> Occurs
occursIn :: forall a. HasFree a => Int -> a -> Occurs
occursIn Int
i a
x = forall a. a -> Maybe a -> a
fromMaybe forall a. Monoid a => a
mempty forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Int
i (forall a. HasFree a => a -> Map Int Occurs
freeVars a
x)

instance HasFree Int where
  freeVars :: Int -> Map Int Occurs
freeVars Int
x = forall k a. k -> a -> Map k a
Map.singleton Int
x Occurs
once

instance HasFree a => HasFree [a] where
  freeVars :: [a] -> Map Int Occurs
freeVars [a]
xs = forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
(a -> a -> a) -> f (Map k a) -> Map k a
Map.unionsWith forall a. Monoid a => a -> a -> a
mappend forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a. HasFree a => a -> Map Int Occurs
freeVars [a]
xs

instance (HasFree a, HasFree b) => HasFree (a, b) where
  freeVars :: (a, b) -> Map Int Occurs
freeVars (a
x, b
y) = forall k a. Ord k => (a -> a -> a) -> Map k a -> Map k a -> Map k a
Map.unionWith forall a. Monoid a => a -> a -> a
mappend (forall a. HasFree a => a -> Map Int Occurs
freeVars a
x) (forall a. HasFree a => a -> Map Int Occurs
freeVars b
y)

data Binder a = Binder Int a

instance HasFree a => HasFree (Binder a) where
  freeVars :: Binder a -> Map Int Occurs
freeVars (Binder Int
0 a
x) = forall a. HasFree a => a -> Map Int Occurs
freeVars a
x
  freeVars (Binder Int
k a
x) = forall k a. (k -> a -> Bool) -> Map k a -> Map k a
Map.filterWithKey (\ Int
k Occurs
_ -> Int
k forall a. Ord a => a -> a -> Bool
>= Int
0) forall a b. (a -> b) -> a -> b
$ forall k1 k2 a. (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeysMonotonic (forall a. Num a => a -> a -> a
subtract Int
k) forall a b. (a -> b) -> a -> b
$ forall a. HasFree a => a -> Map Int Occurs
freeVars a
x

newtype InSeq a = InSeq a

instance HasFree a => HasFree (InSeq a) where
  freeVars :: InSeq a -> Map Int Occurs
freeVars (InSeq a
x) = Occurs -> Occurs
inSeq forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. HasFree a => a -> Map Int Occurs
freeVars a
x

instance HasFree TTerm where
  freeVars :: TTerm -> Map Int Occurs
freeVars = \case
    TDef{}    -> forall k a. Map k a
Map.empty
    TLit{}    -> forall k a. Map k a
Map.empty
    TCon{}    -> forall k a. Map k a
Map.empty
    TPrim{}   -> forall k a. Map k a
Map.empty
    TUnit{}   -> forall k a. Map k a
Map.empty
    TSort{}   -> forall k a. Map k a
Map.empty
    TErased{} -> forall k a. Map k a
Map.empty
    TError{}  -> forall k a. Map k a
Map.empty
    TVar Int
i         -> forall a. HasFree a => a -> Map Int Occurs
freeVars Int
i
    TApp (TPrim TPrim
PSeq) [TVar Int
x, TTerm
b] -> forall a. HasFree a => a -> Map Int Occurs
freeVars (forall a. a -> InSeq a
InSeq Int
x, TTerm
b)
    TApp TTerm
f [TTerm]
ts      -> forall a. HasFree a => a -> Map Int Occurs
freeVars (TTerm
f, [TTerm]
ts)
    TLam TTerm
b         -> Occurs -> Occurs
underLambda forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. HasFree a => a -> Map Int Occurs
freeVars (forall a. Int -> a -> Binder a
Binder Int
1 TTerm
b)
    TLet TTerm
e TTerm
b       -> forall a. HasFree a => a -> Map Int Occurs
freeVars (TTerm
e, forall a. Int -> a -> Binder a
Binder Int
1 TTerm
b)
    TCase Int
i CaseInfo
_ TTerm
d [TAlt]
bs -> forall a. HasFree a => a -> Map Int Occurs
freeVars (Int
i, (TTerm
d, [TAlt]
bs))
    TCoerce TTerm
t      -> forall a. HasFree a => a -> Map Int Occurs
freeVars TTerm
t

instance HasFree TAlt where
  freeVars :: TAlt -> Map Int Occurs
freeVars = \case
    TACon QName
_ Int
i TTerm
b -> forall a. HasFree a => a -> Map Int Occurs
freeVars (forall a. Int -> a -> Binder a
Binder Int
i TTerm
b)
    TALit Literal
_ TTerm
b   -> forall a. HasFree a => a -> Map Int Occurs
freeVars TTerm
b
    TAGuard TTerm
g TTerm
b -> forall a. HasFree a => a -> Map Int Occurs
freeVars (TTerm
g, TTerm
b)

-- | Strenghtening.
tryStrengthen :: (HasFree a, Subst a) => Int -> a -> Maybe a
tryStrengthen :: forall a. (HasFree a, Subst a) => Int -> a -> Maybe a
tryStrengthen Int
n a
t =
  case forall k a. Map k a -> Maybe ((k, a), Map k a)
Map.minViewWithKey (forall a. HasFree a => a -> Map Int Occurs
freeVars a
t) of
    Just ((Int
i, Occurs
_), Map Int Occurs
_) | Int
i forall a. Ord a => a -> a -> Bool
< Int
n -> forall a. Maybe a
Nothing
    Maybe ((Int, Occurs), Map Int Occurs)
_ -> forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall a. Subst a => Substitution' (SubstArg a) -> a -> a
applySubst (forall a. Impossible -> Int -> Substitution' a
strengthenS HasCallStack => Impossible
impossible Int
n) a
t