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

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

import Agda.TypeChecking.Monad

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

    -- Coerce a term `t`. The result (when translated to Haskell) has type
    -- `forall a. a`.
    coerce :: TTerm -> TCM TTerm
coerce TTerm
t =
      case TTerm
t of
        TVar{}    -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> TCM TTerm) -> TTerm -> TCM TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TPrim{}   -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> TCM TTerm) -> TTerm -> TCM TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TDef{}    -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> TCM TTerm) -> TTerm -> TCM TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TCon{}    -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> TCM TTerm) -> TTerm -> TCM TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TLit{}    -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> TCM TTerm) -> TTerm -> TCM TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TUnit{}   -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> TCM TTerm) -> TTerm -> TCM TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TSort{}   -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> TCM TTerm) -> TTerm -> TCM TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TErased{} -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TCoerce{} -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TError{}  -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TApp TTerm
f Args
vs -> do
          Nat
ar <- TTerm -> TCM 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) -> TCMT IO Args -> TCM TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> TCM TTerm) -> Args -> TCMT IO Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> TCM 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) -> TCMT IO Args -> TCM TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> TCM TTerm) -> Args -> TCMT IO Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> TCM 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) -> TCM TTerm -> TCM TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCM TTerm
softCoerce TTerm
b
        TLet TTerm
e TTerm
b       -> TTerm -> TTerm -> TTerm
TLet (TTerm -> TTerm -> TTerm) -> TCM TTerm -> TCMT IO (TTerm -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCM TTerm
softCoerce TTerm
e TCMT IO (TTerm -> TTerm) -> TCM TTerm -> TCM TTerm
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> TCM 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)
-> TCM TTerm -> TCMT IO ([TAlt] -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCM TTerm
coerce TTerm
d TCMT IO ([TAlt] -> TTerm) -> TCMT IO [TAlt] -> TCM TTerm
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TAlt -> TCMT IO TAlt) -> [TAlt] -> TCMT IO [TAlt]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TAlt -> TCMT IO TAlt
coerceAlt [TAlt]
bs

    coerceAlt :: TAlt -> TCMT IO TAlt
coerceAlt (TACon QName
c Nat
a TTerm
b) = QName -> Nat -> TTerm -> TAlt
TACon QName
c Nat
a (TTerm -> TAlt) -> TCM TTerm -> TCMT IO TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCM TTerm
coerce TTerm
b
    coerceAlt (TAGuard TTerm
g TTerm
b) = TTerm -> TTerm -> TAlt
TAGuard   (TTerm -> TTerm -> TAlt) -> TCM TTerm -> TCMT IO (TTerm -> TAlt)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCM TTerm
coerce TTerm
g TCMT IO (TTerm -> TAlt) -> TCM TTerm -> TCMT IO TAlt
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> TCM TTerm
coerce TTerm
b
    coerceAlt (TALit Literal
l TTerm
b)   = Literal -> TTerm -> TAlt
TALit Literal
l   (TTerm -> TAlt) -> TCM TTerm -> TCMT IO TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCM 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 -> TCM TTerm
softCoerce TTerm
t =
      case TTerm
t of
        TVar{}    -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TPrim{}   -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TDef{}    -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TCon{}    -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TLit{}    -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TUnit{}   -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TSort{}   -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TErased{} -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TCoerce{} -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TError{}  -> TTerm -> TCM TTerm
forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TApp TTerm
f Args
vs -> do
          Nat
ar <- TTerm -> TCM 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) -> TCMT IO Args -> TCM TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> TCM TTerm) -> Args -> TCMT IO Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> TCM TTerm
softCoerce Args
vs
            else TTerm -> Args -> TTerm
TApp TTerm
f (Args -> TTerm) -> TCMT IO Args -> TCM TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (TTerm -> TCM TTerm) -> Args -> TCMT IO Args
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TTerm -> TCM TTerm
coerce Args
vs
        TLam TTerm
b         -> TTerm -> TTerm
TLam (TTerm -> TTerm) -> TCM TTerm -> TCM TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCM TTerm
softCoerce TTerm
b
        TLet TTerm
e TTerm
b       -> TTerm -> TTerm -> TTerm
TLet (TTerm -> TTerm -> TTerm) -> TCM TTerm -> TCMT IO (TTerm -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCM TTerm
softCoerce TTerm
e TCMT IO (TTerm -> TTerm) -> TCM TTerm -> TCM TTerm
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> TCM 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)
-> TCM TTerm -> TCMT IO ([TAlt] -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> TCM TTerm
coerce TTerm
d TCMT IO ([TAlt] -> TTerm) -> TCMT IO [TAlt] -> TCM TTerm
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TAlt -> TCMT IO TAlt) -> [TAlt] -> TCMT IO [TAlt]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM TAlt -> TCMT IO TAlt
coerceAlt [TAlt]
bs

funArity :: TTerm -> TCM Nat
funArity :: TTerm -> TCM 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) -> TCMT IO (Maybe TTerm) -> TCM Nat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QName -> TCMT IO (Maybe TTerm)
getTreeless QName
q
funArity (TCon QName
q)  = QName -> TCM Nat
erasedArity QName
q
funArity (TPrim TPrim
_) = Nat -> TCM Nat
forall (m :: * -> *) a. Monad m => a -> m a
return Nat
3 -- max arity of any primitive
funArity TTerm
_         = Nat -> TCM Nat
forall (m :: * -> *) a. Monad m => a -> m a
return Nat
0

-- | The number of retained arguments after erasure.
erasedArity :: QName -> TCM Nat
erasedArity :: QName -> TCM 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) -> TCMT IO [Bool] -> TCM Nat
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> QName -> TCMT IO [Bool]
getErasedConArgs QName
q