module Agda.Compiler.Treeless.AsPatterns (recoverAsPatterns) where

import Control.Monad.Reader

import Agda.Syntax.Treeless

data AsPat = AsPat Int QName [Int]  -- x@(c ys)
  deriving (Int -> AsPat -> ShowS
[AsPat] -> ShowS
AsPat -> String
(Int -> AsPat -> ShowS)
-> (AsPat -> String) -> ([AsPat] -> ShowS) -> Show AsPat
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> AsPat -> ShowS
showsPrec :: Int -> AsPat -> ShowS
$cshow :: AsPat -> String
show :: AsPat -> String
$cshowList :: [AsPat] -> ShowS
showList :: [AsPat] -> ShowS
Show)

wk :: Int -> AsPat -> AsPat
wk :: Int -> AsPat -> AsPat
wk Int
n (AsPat Int
x QName
c [Int]
ys) = Int -> QName -> [Int] -> AsPat
AsPat (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x) QName
c ((Int -> Int) -> [Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+) [Int]
ys)

type S = Reader [AsPat]

runS :: S a -> a
runS :: forall a. S a -> a
runS S a
m = S a -> [AsPat] -> a
forall r a. Reader r a -> r -> a
runReader S a
m []

underBinds :: Int -> S a -> S a
underBinds :: forall a. Int -> S a -> S a
underBinds Int
0 = S a -> S a
forall a. a -> a
id
underBinds Int
n = ([AsPat] -> [AsPat]) -> S a -> S a
forall a.
([AsPat] -> [AsPat])
-> ReaderT [AsPat] Identity a -> ReaderT [AsPat] Identity a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local ((AsPat -> AsPat) -> [AsPat] -> [AsPat]
forall a b. (a -> b) -> [a] -> [b]
map ((AsPat -> AsPat) -> [AsPat] -> [AsPat])
-> (AsPat -> AsPat) -> [AsPat] -> [AsPat]
forall a b. (a -> b) -> a -> b
$ Int -> AsPat -> AsPat
wk Int
n)

bindAsPattern :: AsPat -> S a -> S a
bindAsPattern :: forall a. AsPat -> S a -> S a
bindAsPattern AsPat
p = ([AsPat] -> [AsPat])
-> ReaderT [AsPat] Identity a -> ReaderT [AsPat] Identity a
forall a.
([AsPat] -> [AsPat])
-> ReaderT [AsPat] Identity a -> ReaderT [AsPat] Identity a
forall r (m :: * -> *) a. MonadReader r m => (r -> r) -> m a -> m a
local (AsPat
p AsPat -> [AsPat] -> [AsPat]
forall a. a -> [a] -> [a]
:)

lookupAsPattern :: QName -> [TTerm] -> S TTerm
lookupAsPattern :: QName -> [TTerm] -> S TTerm
lookupAsPattern QName
c [TTerm]
vs
  | Just [Int]
xs <- [TTerm] -> Maybe [Int]
allVars [TTerm]
vs = do
    [AsPat]
ps <- ReaderT [AsPat] Identity [AsPat]
forall r (m :: * -> *). MonadReader r m => m r
ask
    case [ Int
x | AsPat Int
x QName
c' [Int]
ys <- [AsPat]
ps, QName
c QName -> QName -> Bool
forall a. Eq a => a -> a -> Bool
== QName
c', [Int]
xs [Int] -> [Int] -> Bool
forall a. Eq a => a -> a -> Bool
== [Int]
ys ] of
      Int
x : [Int]
_ -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TTerm -> S TTerm) -> TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ Int -> TTerm
TVar Int
x
      [Int]
_     -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] 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 (QName -> TTerm
TCon QName
c) [TTerm]
vs
  | Bool
otherwise = TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] 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 (QName -> TTerm
TCon QName
c) [TTerm]
vs
  where
    allVars :: [TTerm] -> Maybe [Int]
allVars = (TTerm -> Maybe Int) -> [TTerm] -> Maybe [Int]
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 -> Maybe Int
getVar
    getVar :: TTerm -> Maybe Int
getVar (TVar Int
x) = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
x
    getVar TTerm
_        = Maybe Int
forall a. Maybe a
Nothing   -- what about erased?

-- | We lose track of @-patterns in the internal syntax. This pass puts them
--   back.
recoverAsPatterns :: Monad m => TTerm -> m TTerm
recoverAsPatterns :: forall (m :: * -> *). Monad m => TTerm -> m TTerm
recoverAsPatterns TTerm
t = TTerm -> m TTerm
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (TTerm -> m TTerm) -> TTerm -> m TTerm
forall a b. (a -> b) -> a -> b
$ S TTerm -> TTerm
forall a. S a -> a
runS (TTerm -> S TTerm
recover TTerm
t)

recover :: TTerm -> S TTerm
recover :: TTerm -> S TTerm
recover TTerm
t =
  case TTerm
t of
    TApp TTerm
f [TTerm]
vs -> do
      TTerm
f  <- TTerm -> S TTerm
recover TTerm
f
      [TTerm]
vs <- (TTerm -> S TTerm) -> [TTerm] -> ReaderT [AsPat] 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
recover [TTerm]
vs
      TTerm -> [TTerm] -> S TTerm
tApp TTerm
f [TTerm]
vs
    TLam TTerm
b -> TTerm -> TTerm
TLam (TTerm -> TTerm) -> S TTerm -> S TTerm
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> S TTerm -> S TTerm
forall a. Int -> S a -> S a
underBinds Int
1 (TTerm -> S TTerm
recover TTerm
b)
    TCon{} -> TTerm -> [TTerm] -> S TTerm
tApp TTerm
t []   -- need to recover nullary constructors as well (to make deep @-patterns work)
    TLet TTerm
v TTerm
b -> TTerm -> TTerm -> TTerm
TLet (TTerm -> TTerm -> TTerm)
-> S TTerm -> ReaderT [AsPat] Identity (TTerm -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm
recover TTerm
v ReaderT [AsPat] Identity (TTerm -> TTerm) -> S TTerm -> S TTerm
forall a b.
ReaderT [AsPat] Identity (a -> b)
-> ReaderT [AsPat] Identity a -> ReaderT [AsPat] Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> S TTerm -> S TTerm
forall a. Int -> S a -> S a
underBinds Int
1 (TTerm -> S TTerm
recover TTerm
b)
    TCase Int
x CaseInfo
ct TTerm
d [TAlt]
bs -> Int -> CaseInfo -> TTerm -> [TAlt] -> TTerm
TCase Int
x CaseInfo
ct (TTerm -> [TAlt] -> TTerm)
-> S TTerm -> ReaderT [AsPat] Identity ([TAlt] -> TTerm)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm
recover TTerm
d ReaderT [AsPat] Identity ([TAlt] -> TTerm)
-> ReaderT [AsPat] Identity [TAlt] -> S TTerm
forall a b.
ReaderT [AsPat] Identity (a -> b)
-> ReaderT [AsPat] Identity a -> ReaderT [AsPat] Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (TAlt -> ReaderT [AsPat] Identity TAlt)
-> [TAlt] -> ReaderT [AsPat] 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 -> TAlt -> ReaderT [AsPat] Identity TAlt
recoverAlt Int
x) [TAlt]
bs
    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
recover TTerm
t
    TLit{}    -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
    TVar{}    -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
    TPrim{}   -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
    TDef{}    -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
    TUnit{}   -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
    TSort{}   -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
    TErased{} -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t
    TError{}  -> TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TTerm
t

recoverAlt :: Int -> TAlt -> S TAlt
recoverAlt :: Int -> TAlt -> ReaderT [AsPat] Identity TAlt
recoverAlt Int
x TAlt
b =
  case TAlt
b of
    TACon QName
c Int
n TTerm
b -> QName -> Int -> TTerm -> TAlt
TACon QName
c Int
n (TTerm -> TAlt) -> S TTerm -> ReaderT [AsPat] 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
underBinds Int
n (AsPat -> S TTerm -> S TTerm
forall a. AsPat -> S a -> S a
bindAsPattern (Int -> QName -> [Int] -> AsPat
AsPat (Int
x Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n) QName
c [Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1, Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
2..Int
0]) (S TTerm -> S TTerm) -> S TTerm -> S TTerm
forall a b. (a -> b) -> a -> b
$ TTerm -> S TTerm
recover TTerm
b)
    TAGuard TTerm
g TTerm
b -> TTerm -> TTerm -> TAlt
TAGuard (TTerm -> TTerm -> TAlt)
-> S TTerm -> ReaderT [AsPat] Identity (TTerm -> TAlt)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm
recover TTerm
g ReaderT [AsPat] Identity (TTerm -> TAlt)
-> S TTerm -> ReaderT [AsPat] Identity TAlt
forall a b.
ReaderT [AsPat] Identity (a -> b)
-> ReaderT [AsPat] Identity a -> ReaderT [AsPat] Identity b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> TTerm -> S TTerm
recover TTerm
b
    TALit Literal
l TTerm
b   -> Literal -> TTerm -> TAlt
TALit Literal
l (TTerm -> TAlt) -> S TTerm -> ReaderT [AsPat] Identity TAlt
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TTerm -> S TTerm
recover TTerm
b

tApp :: TTerm -> [TTerm] -> S TTerm
tApp :: TTerm -> [TTerm] -> S TTerm
tApp (TCon QName
c) [TTerm]
vs = QName -> [TTerm] -> S TTerm
lookupAsPattern QName
c [TTerm]
vs
tApp TTerm
f [TTerm]
vs        = TTerm -> S TTerm
forall a. a -> ReaderT [AsPat] 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 TTerm
f [TTerm]
vs