module Core.Value where

import Core.Info
import Common
import {-# source #-} qualified Core.Syntax as S

----------------------------------------------------------------------------------------------------

data SCast
  = SCastPi PiClosure PiClosure
    -- ^ Canonical cast along Pi types
  | SCastRigid Ty Ty
    -- ^ Rigidly blocked cast (indicates conversion error)
  | SCastFlex Ty Ty
    -- ^ Flexibly blocked cast
  deriving Show

data Spine
  = SId
  | SApp Spine Tm (Lazy Ty) PiDomKind
  | SCast Spine SCast
  deriving Show

spineHasCast :: Spine -> Bool
spineHasCast = \case
  SCast{}       -> True
  SId           -> False
  SApp sp _ _ _ -> spineHasCast sp

data RSpine
  = RSId
  | RSApp Tm (Lazy Ty) PiDomKind RSpine
  | RSCast SCast RSpine
  deriving Show

reverseSpine :: Spine -> RSpine
reverseSpine = go RSId where
  go acc = \case
    SId            -> acc
    SApp sp t a pd -> go (RSApp t a pd acc) sp
    SCast sp cst   -> go (RSCast cst acc) sp

-- | Spines for unfolding don't need type annotations!
data UnfoldSpine
  = USId
  | USApp UnfoldSpine Tm PiDomKind
  | USCast UnfoldSpine SCast
  deriving Show

unfoldSpineHasCast :: UnfoldSpine -> Bool
unfoldSpineHasCast = \case
  USCast{}     -> True
  USId         -> False
  USApp sp _ _ -> unfoldSpineHasCast sp

data RUnfoldSpine
  = RUSId
  | RUSApp Tm PiDomKind RUnfoldSpine
  | RUSCast SCast RUnfoldSpine
  deriving Show

reverseUnfoldSpine :: UnfoldSpine -> RUnfoldSpine
reverseUnfoldSpine = go RUSId where
  go acc = \case
    USId          -> acc
    USApp sp t pd -> go (RUSApp t pd acc) sp
    USCast sp cst -> go (RUSCast cst acc) sp

spineToUnfoldSpine :: Spine -> UnfoldSpine
spineToUnfoldSpine = \case
  SId            -> USId
  SApp sp t _ pd -> USApp (spineToUnfoldSpine sp) t pd
  SCast sp cst   -> USCast (spineToUnfoldSpine sp) cst

----------------------------------------------------------------------------------------------------

data RigidHead
  = RHVar Lvl
  -- ^ Bound variable
  | RHCast Ty Ty Tm
  -- ^ Types must be rigidly distinct and the value
  --   must be canonical. This represents cast along an /error/.
  deriving Show

data MetaEnv = MetaEnv {
    metaEnvEnv          :: Env
  , metaEnvIsWkIdentity :: Bool
         -- ^ If this is true, the env must be a weakening of an identity sub!
         --   We store the env here for structure sharing.
  }
  deriving Show

data MetaHead = MetaHead {
    metaHeadMetaVar :: MetaVar
  , metaHeadEnv     :: MetaEnv
  }

instance Show MetaHead where
  show (MetaHead x y) = show (x, y)

data FlexHead
  = FHMeta MetaHead
  | FHCast Ty Ty Tm
  -- ^ Type conversion is flexibly blocked, value must be
  --   canonical.
  deriving Show

data Env
  = ENil
  | EExt Env (Lazy Tm) -- ^ Vanilla extension
  | EDef Env (Lazy Tm) -- ^ Definitional extension for local unfolding control.
                       --   NOTE: when elaboration goes under a let, it adds an
                       --   Unfold value here. In other words, the /identity
                       --   environment/ should have an Unfold in each
                       --   EDef. That's more efficient in lookups, where we
                       --   can just literally do a lookup and not build any Unfold.
  deriving Show

type EnvArg = (?env :: Env)

{-# inline wkEnv #-}
wkEnv :: (EnvArg => a) -> (EnvArg => a)
wkEnv act = case ?env of
  EExt e _ -> let ?env = e in act
  _        -> impossible

data PiClosure = PiCl {
    piClosureName    :: Name
  , piClosureDomKind :: PiDomKind
  , piClosureDomTy   :: Lazy Ty
  , piClosureEnv     :: Env
  , piClosureBody    :: S.Ty
  }

instance Show PiClosure where
  show (PiCl x pd a e t) = show (x, pd, a, e, t)

data LamClosureBody
  = LCBEval Env S.Tm
  | LCBCastLam PiClosure PiClosure LamClosureBody
  deriving Show

data LamClosure = LamCl {
    lamClosureName    :: Name
  , lamClosureDomKind :: PiDomKind
  , lamClosureDomTy   :: Lazy Ty
  , lamClosureBody    :: LamClosureBody
  }

instance Show LamClosure where
  show (LamCl x i a t) = show (x, i, a, t)

data LetClosure = LetCl {
    letClosureName :: Name
  , letClosureTy   :: Lazy Ty
  , letClosureEnv  :: Env
  , letClosureBody :: S.Tm
  }

instance Show LetClosure where
  show (LetCl x a e t) = show (x, a, e, t)

data CompareTypes
  = CTPi PiClosure PiClosure
  | CTSame
  | CTDistinct
  | CTFlex
  deriving Show

data UnfoldHead
  = UHMeta MetaHead
  | UHTopDef {-# nounpack #-} DefInfo
  | UHLocalDef Lvl
  deriving Show

--------------------------------------------------------------------------------

-- | Finite levels. The representation is
--     N₀ ⊔ (t₁ + N₁) ⊔ (t₂ + N₂) ⊔ ... ⊔ (tᵢ + Nᵢ)
--   where Nᵢ are natural numberes and tᵢ are values.
--   If tᵢ all have rigid heads and are all rigid distinct, the FinLevelTm is rigid.
--   If there are any flex comparisons between tᵢ or any tᵢ has flex head, the
--   FinLevelTm is flex.
data FinLevelTm = FLT IsFlex Integer FLTTail
  deriving Show

data FLTTail
  = FLTNil
  | FLTCons Tm Integer FLTTail
  deriving Show

instance Sized FLTTail where
  size = go 0 where
    go acc FLTNil             = acc
    go acc (FLTCons _ _ tail) = go (acc + 1) tail

instance Sized FinLevelTm where
  size (FLT _ _ tail) = size tail

pattern FLTSingle :: Tm -> Integer -> FLTTail
pattern FLTSingle t n = FLTCons t n FLTNil

data Level
  = ZeroPlus FinLevelTm
  | OmegaPlus Integer
  | OmegaPlusOmega
  deriving Show

pattern FLTRigid :: Integer -> FLTTail -> FinLevelTm
pattern FLTRigid n tail = FLT False n tail

pattern FLTFlex :: Integer -> FLTTail -> FinLevelTm
pattern FLTFlex n tail = FLT True n tail
{-# complete FLTRigid, FLTFlex #-}

--------------------------------------------------------------------------------

data Relevance
  = Relevant                         -- ^ Definitionally relevant
  | Contractible                     -- ^ Definitionally contractible
  | RelevanceOf Level FlexHead Spine -- ^ Relevance blocked by (Decode size (Flex h sp))
  deriving Show

data RelevanceInfo = RI {
    relevanceInfoTy        :: Ty
  , relevanceInfoRelevance :: Relevance
  }

instance Show RelevanceInfo where
  show (RI a r) = show (a, r)

data Tm
  = Rigid RigidHead Spine (Lazy RelevanceInfo)
  | Flex FlexHead Spine (Lazy Ty)
  | Unfold UnfoldHead UnfoldSpine (Lazy Tm)
  | Lam LamClosure
  | FinLevelTm FinLevelTm
  | Code Level Ty
  deriving Show

data Ty
  = Pi PiClosure
  | FinLevel
  | Level
  | TYPE Level
  | Decode Level Tm
  deriving Show

makeFields ''LamClosure
makeFields ''LetClosure
makeFields ''MetaEnv
makeFields ''MetaHead
makeFields ''PiClosure
makeFields ''RelevanceInfo