{-# LANGUAGE DataKinds #-}
{-# LANGUAGE ExplicitForAll #-}
{-# LANGUAGE TypeOperators #-}
module Language (
  scopeCheck,
  SExpr,
  module Language,
) where

import AST
import Data
import Language.AST


lambda :: forall a args t. KnownTy a => Var a -> SFun args t -> SFun (Append args '[a]) t
lambda var (SFun args e) = SFun (sappend args (var `SCons` SNil)) e

body :: SExpr t -> SFun '[] t
body e = SFun SNil e


data a :-> b = a :-> b
  deriving (Show)
infix 0 :->


TODO
-- TODO: should give SExpr an environment index of kind '[(Symbol, Ty)]. Then
-- the IsLabel instance for SExpr (but not the one for Var!) can check that the
-- type in the named environment matches the locally expected type.

let_ :: KnownTy a => Var a -> SExpr a -> SExpr t -> SExpr t
let_ var rhs e = SELet rhs (Lambda var e)

pair :: SExpr a -> SExpr b -> SExpr (TPair a b)
pair = SEPair

fst_ :: SExpr (TPair a b) -> SExpr a
fst_ = SEFst

snd_ :: SExpr (TPair a b) -> SExpr b
snd_ = SESnd

nil :: SExpr TNil
nil = SENil

inl :: STy b -> SExpr a -> SExpr (TEither a b)
inl = SEInl

inr :: STy a -> SExpr b -> SExpr (TEither a b)
inr = SEInr

case_ :: (KnownTy a, KnownTy b)
      => SExpr (TEither a b) -> (Var a :-> SExpr c) -> (Var b :-> SExpr c) -> SExpr c
case_ e (v1 :-> e1) (v2 :-> e2) = SECase e (Lambda v1 e1) (Lambda v2 e2)

build1 :: SExpr TIx -> (Var TIx :-> SExpr t) -> SExpr (TArr (S Z) t)
build1 a (v :-> b) = SEBuild1 a (Lambda v b)

build :: SNat n -> SExpr (Tup (Replicate n TIx)) -> (Var (Tup (Replicate n TIx)) :-> SExpr t) -> SExpr (TArr n t)
build n a (v :-> b) = SEBuild n a (Lambda v b)

fold1 :: KnownTy t => (SExpr t -> SExpr t -> SExpr t) -> SExpr (TArr (S n) t) -> SExpr (TArr n t)
fold1 f e = SEFold1 (mkLambda2 (f, e) f) e

unit :: SExpr t -> SExpr (TArr Z t)
unit = SEUnit

const_ :: KnownScalTy t => ScalRep t -> SExpr (TScal t)
const_ x =
  let ty = knownScalTy
  in case scalRepIsShow ty of
       Dict -> SEConst ty x

idx0 :: SExpr (TArr Z t) -> SExpr t
idx0 = SEIdx0

(.!) :: SExpr (TArr (S n) t) -> SExpr TIx -> SExpr (TArr n t)
(.!) = SEIdx1

(!) :: SNat n -> SExpr (TArr n t) -> SExpr (Tup (Replicate n TIx)) -> SExpr t
(!) = SEIdx

shape :: SExpr (TArr n t) -> SExpr (Tup (Replicate n TIx))
shape = SEShape

oper :: SOp a t -> SExpr a -> SExpr t
oper = SEOp

error_ :: KnownTy t => String -> SExpr t
error_ s = SEError knownTy s

(.==) :: KnownScalTy st => SExpr (TScal st) -> SExpr (TScal st) -> SExpr (TScal TBool)
a .== b = oper (OEq knownScalTy) (pair a b)

(.<) :: KnownScalTy st => SExpr (TScal st) -> SExpr (TScal st) -> SExpr (TScal TBool)
a .< b = oper (OLt knownScalTy) (pair a b)

(.>) :: KnownScalTy st => SExpr (TScal st) -> SExpr (TScal st) -> SExpr (TScal TBool)
(.>) = flip (.<)

(.<=) :: KnownScalTy st => SExpr (TScal st) -> SExpr (TScal st) -> SExpr (TScal TBool)
a .<= b = oper (OLe knownScalTy) (pair a b)

(.>=) :: KnownScalTy st => SExpr (TScal st) -> SExpr (TScal st) -> SExpr (TScal TBool)
(.>=) = flip (.<=)

not_ :: SExpr (TScal TBool) -> SExpr (TScal TBool)
not_ = oper ONot

if_ :: SExpr (TScal TBool) -> SExpr t -> SExpr t -> SExpr t
if_ e a b = case_ (oper OIf e) (\_ -> a) (\_ -> b)