{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE InstanceSigs #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-partial-fields #-}

module Symantic.Compiler.Term where

import Control.Applicative (Applicative (..))
import Data.Function ((.))
import Data.Function qualified as Fun
import GHC.Types (Constraint, Type)
import Symantic.Semantics.Forall
import Symantic.Syntaxes.Classes (Syntaxes, Unabstractable (..))
import Type.Reflection (Typeable)
import Unsafe.Coerce (unsafeCoerce)

import Symantic.Typer.Type (Ty, monoTy)

type Semantic = Type -> Type
type Syntax = Semantic -> Constraint

-- * Class 'AbstractableTy'
class AbstractableTy meta sem where
  -- | Lambda term abstraction, in HOAS (Higher-Order Abstract Syntax) style.
  lamTy :: Ty meta '[] a -> (sem a -> sem b) -> sem (a -> b)

fun ::
  AbstractableTy () sem =>
  Typeable a =>
  (sem a -> sem b) ->
  sem (a -> b)
fun = lamTy (monoTy @_ @())

-- * Class 'Constable'
class Constable sem where
  constI :: sem (a -> a)
  constK :: sem (a -> b -> a)
  constS :: sem ((a -> b -> c) -> (a -> b) -> a -> c)
  constB :: sem ((b -> c) -> (a -> b) -> a -> c)
  constC :: sem ((a -> b -> c) -> b -> a -> c)

instance (forall sem. Syntaxes syns sem => Constable sem) => Constable (Forall syns) where
  constI = Forall constI
  constK = Forall constK
  constS = Forall constS
  constB = Forall constB
  constC = Forall constC
instance
  ( forall sem. Syntaxes syns sem => AbstractableTy meta sem
  -- , forall sem a. syn sem => AbstractableLam sem a
  -- , forall sem. syn sem => AbstractableLam sem a
  -- , forall sem. syn sem => Typeable sem -- user instance not accepted
  -- , forall s1 s2. (syn s1, syn s2) => s1 ~ s2 -- crazy...
  ) =>
  AbstractableTy meta (Forall syns)
  where
  lamTy aTy f = Forall (lamTy aTy (\a -> let Forall b = f (forallSem a) in b))
    where
      -- Safe here because (a :: sem a) and (b :: sem b) share the same 'sem'.
      forallSem :: sem a -> Forall syns a
      forallSem a = Forall (unsafeCoerce a)

-- * Type 'OpenTerm'
data OpenTerm (syns :: [Syntax]) (vs :: [Type]) (a :: Type) where
  -- E contains embedded closed (i.e. already abstracted) terms.
  E :: Forall syns a -> OpenTerm syns vs a
  -- V represents a reference to the innermost/top environment
  -- variable, i.e. Z
  V :: OpenTerm syns (a ': vs) a
  -- N represents internalizing the innermost bound variable as a
  -- function argument. In other words, we can represent an open
  -- term referring to a certain variable as a function which
  -- takes that variable as an argument.
  N :: OpenTerm syns vs (a -> b) -> OpenTerm syns (a ': vs) b
  -- For efficiency, there is also a special variant of N for the
  -- case where the term does not refer to the topmost variable at all.
  W :: OpenTerm syns vs b -> OpenTerm syns (a ': vs) b
instance
  ( forall sem. Syntaxes syns sem => AbstractableTy meta sem
  , Syntaxes syns (Forall syns)
  ) =>
  AbstractableTy meta (OpenTerm syns '[])
  where
  lamTy aTy f = E (lamTy aTy (runOpenTerm . f . E))
instance
  ( forall sem. Syntaxes syns sem => Constable sem
  , Syntaxes syns (Forall syns)
  ) =>
  Constable (OpenTerm syns vs)
  where
  constI = E constI
  constK = E constK
  constS = E constS
  constB = E constB
  constC = E constC

-- * Type 'R'
newtype R a = R {unR :: a}
instance AbstractableTy meta R where
  lamTy _aTy f = R (unR . f . R)
instance Unabstractable R where
  R f .@ R x = R (f x)
instance Constable R where
  constI = R Fun.id
  constK = R Fun.const
  constS = R (<*>)
  constB = R (.)
  constC = R Fun.flip

runOpenTerm :: OpenTerm syns '[] a -> Forall syns a
runOpenTerm t = case t of E t' -> t'

eval :: Syntaxes syns R => OpenTerm syns '[] a -> a
eval t
  | Forall sem <- runOpenTerm t =
      unR sem

instance
  ( forall sem. Syntaxes syns sem => Constable sem
  , forall sem. Syntaxes syns sem => Unabstractable sem
  , Syntaxes syns (Forall syns)
  ) =>
  Unabstractable (OpenTerm syns vs)
  where
  (.@) = appOpenTerm

appOpenTerm ::
  forall syns as a b.
  ( forall sem. Syntaxes syns sem => Constable sem
  , forall sem. Syntaxes syns sem => Unabstractable sem
  , Syntaxes syns (Forall syns)
  ) =>
  OpenTerm syns as (a -> b) ->
  OpenTerm syns as a ->
  OpenTerm syns as b
W e1 `appOpenTerm` W e2 = W (e1 `appOpenTerm` e2)
W e `appOpenTerm` E d = W (e `appOpenTerm` E d)
E d `appOpenTerm` W e = W (E d `appOpenTerm` e)
W e `appOpenTerm` V = N e
V `appOpenTerm` W e = N (E (constC .@ constI) `appOpenTerm` e)
W e1 `appOpenTerm` N e2 = N (E constB `appOpenTerm` e1 `appOpenTerm` e2)
N e1 `appOpenTerm` W e2 = N (E constC `appOpenTerm` e1 `appOpenTerm` e2)
N e1 `appOpenTerm` N e2 = N (E constS `appOpenTerm` e1 `appOpenTerm` e2)
N e `appOpenTerm` V = N (E constS `appOpenTerm` e `appOpenTerm` E constI)
V `appOpenTerm` N e = N (E (constS .@ constI) `appOpenTerm` e)
E d `appOpenTerm` N e = N (E (constB .@ d) `appOpenTerm` e)
E d `appOpenTerm` V = N (E d)
V `appOpenTerm` E d = N (E (constC .@ constI .@ d))
N e `appOpenTerm` E d = N (E (constC .@ constC .@ d) `appOpenTerm` e)
E d1 `appOpenTerm` E d2 = E (d1 .@ d2)