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 :: forall (m :: * -> *). HasConstInfo m => TTerm -> m TTerm
addCoercions = forall (m :: * -> *). HasConstInfo m => TTerm -> m TTerm
coerceTop
  where
    -- Don't coerce top-level lambdas.
    coerceTop :: TTerm -> f TTerm
coerceTop (TLam TTerm
b) = TTerm -> TTerm
TLam forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> f TTerm
coerceTop TTerm
b
    coerceTop TTerm
t        = forall (m :: * -> *). HasConstInfo m => TTerm -> m 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{}    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TPrim{}   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TDef{}    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TCon{}    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TLit{}    -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TUnit{}   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TSort{}   -> forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ TTerm -> TTerm
TCoerce TTerm
t
        TErased{} -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TCoerce{} -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TError{}  -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TApp TTerm
f Args
vs -> do
          Int
ar <- forall (m :: * -> *). HasConstInfo m => TTerm -> m Int
funArity TTerm
f
          if forall (t :: * -> *) a. Foldable t => t a -> Int
length Args
vs forall a. Ord a => a -> a -> Bool
> Int
ar
            then TTerm -> Args -> TTerm
TApp (TTerm -> TTerm
TCoerce TTerm
f) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. TTerm -> Args -> TTerm
TApp TTerm
f forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 forall b c a. (b -> c) -> (a -> b) -> a -> c
. TTerm -> TTerm
TLam 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
softCoerce TTerm
e forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> m TTerm
coerce TTerm
b
        TCase Int
x CaseInfo
t TTerm
d [TAlt]
bs -> Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
x CaseInfo
t forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
coerce TTerm
d forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> 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 Int
a TTerm
b) = QName -> Int -> TTerm -> TAlt
TACon QName
c Int
a 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   forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
coerce TTerm
g 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   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{}    -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TPrim{}   -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TDef{}    -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TCon{}    -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TLit{}    -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TUnit{}   -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TSort{}   -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TErased{} -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TCoerce{} -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TError{}  -> forall (m :: * -> *) a. Monad m => a -> m a
return TTerm
t
        TApp TTerm
f Args
vs -> do
          Int
ar <- forall (m :: * -> *). HasConstInfo m => TTerm -> m Int
funArity TTerm
f
          if forall (t :: * -> *) a. Foldable t => t a -> Int
length Args
vs forall a. Ord a => a -> a -> Bool
> Int
ar
            then TTerm -> Args -> TTerm
TApp (TTerm -> TTerm
TCoerce TTerm
f) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> 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 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 forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
softCoerce TTerm
e forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> m TTerm
softCoerce TTerm
b
        TCase Int
x CaseInfo
t TTerm
d [TAlt]
bs -> Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
x CaseInfo
t forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> m TTerm
coerce TTerm
d forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> 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 :: forall (m :: * -> *). HasConstInfo m => TTerm -> m Int
funArity (TDef QName
q)  = forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
0 (forall a b. (a, b) -> a
fst forall b c a. (b -> c) -> (a -> b) -> a -> c
. TTerm -> (Int, TTerm)
tLamView) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *). HasConstInfo m => QName -> m (Maybe TTerm)
getTreeless QName
q
funArity (TCon QName
q)  = forall (m :: * -> *). HasConstInfo m => QName -> m Int
erasedArity QName
q
funArity (TPrim TPrim
_) = forall (m :: * -> *) a. Monad m => a -> m a
return Int
3 -- max arity of any primitive
funArity TTerm
_         = forall (m :: * -> *) a. Monad m => a -> m a
return Int
0

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