module Agda.Compiler.MAlonzo.Coerce (addCoercions, erasedArity) where

import Agda.Syntax.Common (Nat)
import Agda.Syntax.Treeless

import Agda.TypeChecking.Monad
  ( HasConstInfo
  , getErasedConArgs
  , getTreeless
  )

-- | Insert unsafeCoerce (in the form of 'TCoerce') everywhere it's needed in
--   the right-hand side of a definition.
addCoercions :: HasConstInfo m => TTerm -> m TTerm
addCoercions :: TTerm -> m TTerm
addCoercions = TTerm -> m TTerm
forall (f :: * -> *). HasConstInfo f => TTerm -> f TTerm
coerceTop
  where
    -- Don't coerce top-level lambdas.
    coerceTop :: TTerm -> f TTerm
coerceTop (TLam TTerm
b) = TTerm -> TTerm
TLam (TTerm -> TTerm) -> f TTerm -> f TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> f TTerm
coerceTop TTerm
b
    coerceTop TTerm
t        = TTerm -> f TTerm
forall (f :: * -> *). HasConstInfo f => TTerm -> f TTerm
coerce TTerm
t

    -- Coerce a term `t`. The result (when translated to Haskell) has type
    -- `forall a. a`.
    coerce :: TTerm -> m TTerm
coerce TTerm
t =
      case TTerm
t of
        TVar{}    -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> m TTerm) -> TTerm -> m TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TPrim{}   -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> m TTerm) -> TTerm -> m TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TDef{}    -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> m TTerm) -> TTerm -> m TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TCon{}    -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> m TTerm) -> TTerm -> m TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TLit{}    -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> m TTerm) -> TTerm -> m TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TUnit{}   -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> m TTerm) -> TTerm -> m TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TSort{}   -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> m TTerm) -> TTerm -> m TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TErased{} -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TCoerce{} -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TError{}  -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TApp TTerm
f Args
vs -> do
          Nat
ar <- TTerm -> m Nat
forall (m :: * -> *). HasConstInfo m => TTerm -> m Nat
funArity TTerm
f
          if Args -> Nat
forall (t :: * -> *) a. Foldable t => t a -> Nat
length Args
vs Nat -> Nat -> Bool
forall a. Ord a => a -> a -> Bool
> Nat
ar
            then TTerm -> Args -> TTerm
TApp (TTerm -> TTerm
TCoerce TTerm
f) (Args -> TTerm) -> m Args -> m TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> m TTerm) -> Args -> m Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> m TTerm
softCoerce Args
vs
            else TTerm -> TTerm
TCoerce (TTerm -> TTerm) -> (Args -> TTerm) -> Args -> TTerm
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TTerm -> Args -> TTerm
TApp TTerm
f (Args -> TTerm) -> m Args -> m TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> m TTerm) -> Args -> m Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> m TTerm
coerce Args
vs
        TLam TTerm
b         -> TTerm -> TTerm
TCoerce (TTerm -> TTerm) -> (TTerm -> TTerm) -> TTerm -> TTerm
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TTerm -> TTerm
TLam (TTerm -> TTerm) -> m TTerm -> m TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
softCoerce TTerm
b
        TLet TTerm
e TTerm
b       -> TTerm -> TTerm -> TTerm
TLet (TTerm -> TTerm -> TTerm) -> m TTerm -> m (TTerm -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
softCoerce TTerm
e m (TTerm -> TTerm) -> m TTerm -> m TTerm
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> m TTerm
coerce TTerm
b
        TCase Nat
x CaseInfo
t TTerm
d [TAlt]
bs -> Nat -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Nat
x CaseInfo
t (TTerm -> [TAlt] -> TTerm) -> m TTerm -> m ([TAlt] -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
coerce TTerm
d m ([TAlt] -> TTerm) -> m [TAlt] -> m TTerm
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TAlt -> m TAlt) -> [TAlt] -> m [TAlt]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TAlt -> m TAlt
coerceAlt [TAlt]
bs

    coerceAlt :: TAlt -> m TAlt
coerceAlt (TACon QName
c Nat
a TTerm
b) = QName -> Nat -> TTerm -> TAlt
TACon QName
c Nat
a (TTerm -> TAlt) -> m TTerm -> m TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
coerce TTerm
b
    coerceAlt (TAGuard TTerm
g TTerm
b) = TTerm -> TTerm -> TAlt
TAGuard   (TTerm -> TTerm -> TAlt) -> m TTerm -> m (TTerm -> TAlt)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
coerce TTerm
g m (TTerm -> TAlt) -> m TTerm -> m TAlt
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> m TTerm
coerce TTerm
b
    coerceAlt (TALit Literal
l TTerm
b)   = Literal -> TTerm -> TAlt
TALit Literal
l   (TTerm -> TAlt) -> m TTerm -> m TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
coerce TTerm
b

    -- Insert TCoerce in subterms. When translated to Haskell, the resulting
    -- term is well-typed with some type arbitrary type.
    softCoerce :: TTerm -> m TTerm
softCoerce TTerm
t =
      case TTerm
t of
        TVar{}    -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TPrim{}   -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TDef{}    -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TCon{}    -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TLit{}    -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TUnit{}   -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TSort{}   -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TErased{} -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TCoerce{} -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TError{}  -> TTerm -> m TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TApp TTerm
f Args
vs -> do
          Nat
ar <- TTerm -> m Nat
forall (m :: * -> *). HasConstInfo m => TTerm -> m Nat
funArity TTerm
f
          if Args -> Nat
forall (t :: * -> *) a. Foldable t => t a -> Nat
length Args
vs Nat -> Nat -> Bool
forall a. Ord a => a -> a -> Bool
> Nat
ar
            then TTerm -> Args -> TTerm
TApp (TTerm -> TTerm
TCoerce TTerm
f) (Args -> TTerm) -> m Args -> m TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> m TTerm) -> Args -> m Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> m TTerm
softCoerce Args
vs
            else TTerm -> Args -> TTerm
TApp TTerm
f (Args -> TTerm) -> m Args -> m TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> m TTerm) -> Args -> m Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> m TTerm
coerce Args
vs
        TLam TTerm
b         -> TTerm -> TTerm
TLam (TTerm -> TTerm) -> m TTerm -> m TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
softCoerce TTerm
b
        TLet TTerm
e TTerm
b       -> TTerm -> TTerm -> TTerm
TLet (TTerm -> TTerm -> TTerm) -> m TTerm -> m (TTerm -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
softCoerce TTerm
e m (TTerm -> TTerm) -> m TTerm -> m TTerm
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> m TTerm
softCoerce TTerm
b
        TCase Nat
x CaseInfo
t TTerm
d [TAlt]
bs -> Nat -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Nat
x CaseInfo
t (TTerm -> [TAlt] -> TTerm) -> m TTerm -> m ([TAlt] -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
coerce TTerm
d m ([TAlt] -> TTerm) -> m [TAlt] -> m TTerm
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TAlt -> m TAlt) -> [TAlt] -> m [TAlt]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TAlt -> m TAlt
coerceAlt [TAlt]
bs

funArity :: HasConstInfo m => TTerm -> m Nat
funArity :: TTerm -> m Nat
funArity (TDef QName
q)  = Nat -> (TTerm -> Nat) -> Maybe TTerm -> Nat
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Nat
0 ((Nat, TTerm) -> Nat
forall a b. (a, b) -> a
fst ((Nat, TTerm) -> Nat) -> (TTerm -> (Nat, TTerm)) -> TTerm -> Nat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TTerm -> (Nat, TTerm)
tLamView) (Maybe TTerm -> Nat) -> m (Maybe TTerm) -> m Nat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QName -> m (Maybe TTerm)
forall (m :: * -> *). HasConstInfo m => QName -> m (Maybe TTerm)
getTreeless QName
q
funArity (TCon QName
q)  = QName -> m Nat
forall (m :: * -> *). HasConstInfo m => QName -> m Nat
erasedArity QName
q
funArity (TPrim TPrim
_) = Nat -> m Nat
forall (m :: * -> *) a. Monad m => a -> m a
return Nat
3 -- max arity of any primitive
funArity TTerm
_         = Nat -> m Nat
forall (m :: * -> *) a. Monad m => a -> m a
return Nat
0

-- | The number of retained arguments after erasure.
erasedArity :: HasConstInfo m => QName -> m Nat
erasedArity :: QName -> m Nat
erasedArity QName
q = [Bool] -> Nat
forall (t :: * -> *) a. Foldable t => t a -> Nat
length ([Bool] -> Nat) -> ([Bool] -> [Bool]) -> [Bool] -> Nat
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Bool -> Bool) -> [Bool] -> [Bool]
forall a. (a -> Bool) -> [a] -> [a]
filter Bool -> Bool
not ([Bool] -> Nat) -> m [Bool] -> m Nat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QName -> m [Bool]
forall (m :: * -> *). HasConstInfo m => QName -> m [Bool]
getErasedConArgs QName
q