{-# LANGUAGE OverloadedStrings #-}

-- | Substitute AST nodes with magical primitives.
module IR.SubstMagic (
  substMagic,
) where

import qualified IR.IR as I

import Data.Data (Proxy (..))
import Data.Generics (
  Data (..),
  everywhere,
  mkT,
 )


{- | Substitute AST nodes with magical primitives.

Implemented as a syb-style generic tree traversal.

Example usage:

@@
-- Given:
myExpr :: Expr Poly.Type

-- Do:
substMagic (Proxy :: Proxy Poly.Type) myExpr

-- Given:
myProm :: Program Annotated.Type

-- Do:
substMagic (Proxy :: Proxy Annotated.Type) myProg
@@
-}
substMagic :: (Data t, Data a) => Proxy t -> a -> a
substMagic :: Proxy t -> a -> a
substMagic Proxy t
p = (forall a. Data a => a -> a) -> forall a. Data a => a -> a
everywhere ((forall a. Data a => a -> a) -> forall a. Data a => a -> a)
-> (forall a. Data a => a -> a) -> forall a. Data a => a -> a
forall a b. (a -> b) -> a -> b
$ (Expr t -> Expr t) -> a -> a
forall a b. (Typeable a, Typeable b) => (b -> b) -> a -> a
mkT ((Expr t -> Expr t) -> a -> a) -> (Expr t -> Expr t) -> a -> a
forall a b. (a -> b) -> a -> b
$ Proxy t -> Expr t -> Expr t
forall t. Proxy t -> Expr t -> Expr t
substMagicExpr Proxy t
p


-- | Replace applications to built-in names with corresponding primitives.
substMagicExpr :: Proxy t -> I.Expr t -> I.Expr t
substMagicExpr :: Proxy t -> Expr t -> Expr t
substMagicExpr Proxy t
_ Expr t
e = case Expr t -> (Expr t, [(Expr t, t)])
forall t. Expr t -> (Expr t, [(Expr t, t)])
I.unfoldApp Expr t
e of
  (I.Var VarId
"new" t
_, (Expr t
a, t
t) : [(Expr t, t)]
ats) -> Expr t -> [(Expr t, t)] -> Expr t
forall t. Expr t -> [(Expr t, t)] -> Expr t
I.foldApp (Primitive -> [Expr t] -> t -> Expr t
forall t. Primitive -> [Expr t] -> t -> Expr t
I.Prim Primitive
I.New [Expr t
a] t
t) [(Expr t, t)]
ats
  (I.Var VarId
"dup" t
_, (Expr t
a, t
t) : [(Expr t, t)]
ats) -> Expr t -> [(Expr t, t)] -> Expr t
forall t. Expr t -> [(Expr t, t)] -> Expr t
I.foldApp (Primitive -> [Expr t] -> t -> Expr t
forall t. Primitive -> [Expr t] -> t -> Expr t
I.Prim Primitive
I.Dup [Expr t
a] t
t) [(Expr t, t)]
ats
  (I.Var VarId
"drop" t
_, (Expr t
a, t
t) : (Expr t
b, t
_) : [(Expr t, t)]
ats) ->
    Expr t -> [(Expr t, t)] -> Expr t
forall t. Expr t -> [(Expr t, t)] -> Expr t
I.foldApp (Primitive -> [Expr t] -> t -> Expr t
forall t. Primitive -> [Expr t] -> t -> Expr t
I.Prim Primitive
I.Drop [Expr t
a, Expr t
b] t
t) [(Expr t, t)]
ats
  (Expr t, [(Expr t, t)])
_ -> Expr t
e