{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE TypeSynonymInstances  #-}

module AST where

import Base
import TEq
import Repr

-- * The AST type from the paper

infix 1 :=

data Expr   =  Const  Int
            |  Add    Expr  Expr
            |  Mul    Expr  Expr
            |  EVar   Var
            |  Let    Decl  Expr
  deriving Show

data Decl   =  Var := Expr
            |  Seq    Decl  Decl
  deriving Show

type Var   =  String

-- * Instantiating the library for AST 

data AST :: * -> * where
  Expr  ::  AST Expr
  Decl  ::  AST Decl
  Var   ::  AST Var

type instance PF AST  =    
                   K Int                  :>:  Expr  -- |Const|
              :+:  (I Expr :*: I Expr)    :>:  Expr  -- |Add|
              :+:  (I Expr :*: I Expr)    :>:  Expr  -- |Mul|
              :+:  I Var                  :>:  Expr  -- |EVar|
              :+:  (I Decl :*: I Expr)    :>:  Expr  -- |Let|
              :+:  (I Var :*: I Expr)     :>:  Decl  -- |:=|
              :+:  (I Decl :*: I Decl)    :>:  Decl  -- |Seq|
              :+:  K String               :>:  Var   -- |V|

instance Cons AST where
  cons _ = [ "Const", "Add", "Mul", "Let", ":=", "Seq", "V" ]

instance Ix AST Expr where

  from (Const i)   =  L                     (Tag (K i))
  from (Add e f)   =  R (L                  (Tag (I (I0 e) :*: I (I0 f))))
  from (Mul e f)   =  R (R (L               (Tag (I (I0 e) :*: I (I0 f)))))
  from (EVar x)    =  R (R (R (L            (Tag (I (I0 x))))))
  from (Let d e)   =  R (R (R (R (L         (Tag (I (I0 d) :*: I (I0 e)))))))

  to (L                     (Tag (K i)))                        =  Const i
  to (R (L                  (Tag (I (I0 e) :*: I (I0 f)))))     =  Add e f
  to (R (R (L               (Tag (I (I0 e) :*: I (I0 f))))))    =  Mul e f
  to (R (R (R (L            (Tag (I (I0 x)))))))                =  EVar x
  to (R (R (R (R (L         (Tag (I (I0 d) :*: I (I0 e))))))))  =  Let d e
  
  index  =  Expr

instance Ix AST Decl where

  from (x := e)    =  R (R (R (R (R (L      (Tag (I (I0 x) :*: I (I0 e))))))))
  from (Seq c d)   =  R (R (R (R (R (R (L   (Tag (I (I0 c) :*: I (I0 d)))))))))

  to (R (R (R (R (R (L      (Tag (I (I0 x) :*: I (I0 e))))))))) =  x := e
  to (R (R (R (R (R (R (L   (Tag (I (I0 c) :*: I (I0 d))))))))))=  Seq c d

  index  =  Decl

instance Ix AST Var where

  from x           =  R (R (R (R (R (R (R(Tag (K x))))))))

  to (R (R (R (R (R (R (R(Tag (K x)))))))))           =  x

  index  =  Var

-- * Runtime comparison of representations.

instance Repr AST where
  ixeq = ast_eq

ast_eq :: Monad m => AST ix -> AST ix' -> m (ix :=: ix')
ast_eq Expr Expr = return Refl
ast_eq Decl Decl = return Refl
ast_eq Var  Var  = return Refl
ast_eq _    _    = fail "Failed to match witnesses of ED datatypes"