{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module CHAD.Top where

import AST
import CHAD
import CHAD.Types
import Data


type family MergeEnv env where
  MergeEnv '[] = '[]
  MergeEnv (t : ts) = "merge" : MergeEnv ts

mergeDescr :: SList STy env -> Descr env (MergeEnv env)
mergeDescr SNil = DTop
mergeDescr (t `SCons` env) = mergeDescr env `DPush` (t, SMerge)

mergeEnvNoAccum :: SList f env -> Select env (MergeEnv env) "accum" :~: '[]
mergeEnvNoAccum SNil = Refl
mergeEnvNoAccum (_ `SCons` env) | Refl <- mergeEnvNoAccum env = Refl

mergeEnvOnlyMerge :: SList f env -> Select env (MergeEnv env) "merge" :~: env
mergeEnvOnlyMerge SNil = Refl
mergeEnvOnlyMerge (_ `SCons` env) | Refl <- mergeEnvOnlyMerge env = Refl

d1Identity :: STy t -> D1 t :~: t
d1Identity = \case
  STNil -> Refl
  STPair a b | Refl <- d1Identity a, Refl <- d1Identity b -> Refl
  STEither a b | Refl <- d1Identity a, Refl <- d1Identity b -> Refl
  STMaybe t | Refl <- d1Identity t -> Refl
  STArr _ t | Refl <- d1Identity t -> Refl
  STScal _ -> Refl
  STAccum{} -> error "Accumulators not allowed in input program"

d1eIdentity :: SList STy env -> D1E env :~: env
d1eIdentity SNil = Refl
d1eIdentity (t `SCons` env) | Refl <- d1Identity t, Refl <- d1eIdentity env = Refl

chad :: SList STy env -> Ex env t -> Ex (D2 t : D1E env) (TPair (D1 t) (Tup (D2E env)))
chad env term
  | Refl <- mergeEnvNoAccum env
  , Refl <- mergeEnvOnlyMerge env
  = freezeRet (mergeDescr env) (drev (mergeDescr env) term)

chad' :: SList STy env -> Ex env t -> Ex (D2 t : env) (TPair t (Tup (D2E env)))
chad' env term
  | Refl <- d1eIdentity env, Refl <- d1Identity (typeOf term)
  = chad env term