path: root/src/Language/Fiddle/Ast.hs

{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DefaultSignatures #-}
{-# LANGUAGE DeriveAnyClass #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE IncoherentInstances #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Language.Fiddle.Ast where

import Data.Functor.Identity
import Data.Kind (Type)
import Data.List.NonEmpty
import Data.Proxy
import Data.Text (Text)
import Data.Traversable
import Data.Typeable
import GHC.Generics
import GHC.TypeLits

-- The Kind types should have for SyntaxTrees.
--
-- A synatx tree kind should take as parameters, an applicative functor for the
-- first argument. This applicative functor allows errors and warnings to be
-- contained to their deferred regions during compilation, so more errors may be
-- reported an once.
--
-- The second argument is the annotation type. Every syntax tree element has an
-- annotation. When initially parsing, this annotation is a list of comments
-- that came before the element and the source position of the element.
type SynTreeKind = (Type -> Type) -> Type -> Type

-- The Type of number during each stage of compilation. When in the first stage,
-- numbers are just strings like anything else. In later stages, numbers get
-- parsed into actual integers. This makes it easier to process later.
type family NumberType (a :: Stage) :: Type where
  NumberType Stage1 = Text
  NumberType Stage2 = Integer
  NumberType Stage3 = Integer

-- The type that represents an import statement. In the early stages of
-- compilation, this is just a string representing the import path, but in later
-- stages of compilation, this actually gets replaced by an abstract
-- representation of the imported material.
type family ImportType (stage :: Stage) :: SynTreeKind where
  ImportType Stage1 = ImportStatement
  ImportType Stage2 = ImportStatement
  ImportType Stage3 = ImportStatement

-- Type-level constraint to determine if a stage is less than some natural
-- ordinal. Used to bound parts of the AST in multiple stages.
type StageLessThan stage (n :: Natural) = (CmpNat (StageNumber stage) n ~ LT)

-- A Name is multiple identifiers separated by dots. It's the way of namespaces
-- to different packages.
data Name f a where
  Name :: NonEmpty (Identifier f a) -> a -> Name f a
  deriving (Generic, Annotated, Alter, Typeable)

-- [[packed, rust: name="field_name"]]
data Directive f a where
  Directive :: f (DirectiveBody f a) -> a -> Directive f a
  deriving (Generic, Annotated, Alter, Typeable)

data DirectiveBody f a where
  DirectiveBody :: [DirectiveElement f a] -> a -> DirectiveBody f a
  deriving (Generic, Annotated, Alter, Typeable)

data DirectiveElement f a where
  -- <key>
  DirectiveElementKey ::
    Maybe (Identifier f a) -> Identifier f a -> a -> DirectiveElement f a
  -- (<backend>:)? <key>=<value>
  DirectiveElementKeyValue ::
    Maybe (Identifier f a) ->
    Identifier f a ->
    DirectiveExpression f a ->
    a ->
    DirectiveElement f a
  deriving (Generic, Annotated, Alter, Typeable)

data DirectiveExpression f a where
  DirectiveString :: Text -> a -> DirectiveExpression f a
  DirectiveNumber :: Text -> a -> DirectiveExpression f a
  deriving (Generic, Annotated, Alter, Typeable)

data Directed t f a where
  Directed :: [Directive f a] -> t f a -> a -> Directed t f a
  deriving (Generic, Annotated, Alter, Typeable)

mapDirected :: (t f a -> t' f a) -> Directed t f a -> Directed t' f a
mapDirected fn (Directed dr tfa a) = Directed dr (fn tfa) a

mapDirectedM ::
  (Monad m) => (t f a -> m (t' f a)) -> Directed t f a -> m (Directed t' f a)
mapDirectedM fn (Directed dr tfa a) = Directed dr <$> fn tfa <*> pure a

asDirected :: (Annotated t) => t f a -> Directed t f a
asDirected tfa = Directed [] tfa (annot tfa)

undirected :: Directed t f a -> t f a
undirected (Directed _ tfa _) = tfa

-- Stage of compilation. Parts of the AST maybe un unavailable with other stages
-- as compilation simplifies the AST.
data Stage = Stage1 | Stage2 | Stage3
  deriving (Typeable)

-- Root of the parse tree. Just contains a list of declarations.
data FiddleUnit (stage :: Stage) (f :: Type -> Type) a where
  FiddleUnit :: [Directed (FiddleDecl stage) f a] -> a -> FiddleUnit stage f a
  deriving (Generic, Annotated, Typeable)

deriving instance (Alter (ImportType stage)) => Alter (FiddleUnit stage)

-- Just an identifier.
data Identifier f a = Identifier !Text a
  deriving (Generic, Annotated, Alter, Typeable)

-- Expression.
data Expression stage f a where
  -- Just a string. Parsing the number comes in stage2.
  LitNum :: NumberType stage -> a -> Expression stage f a
  Var :: Identifier f a -> a -> Expression stage f a

-- Wrapper class because using a type family like "ImportType" means I cannot
-- derive classes. I'd rather only derive for this class than have to derive for
-- the whole PackageBody class.
data ImportTypeWrapper stage f a where
  ImportTypeWrapper :: ImportType stage f a -> ImportTypeWrapper stage f a

instance (Alter (ImportType stage)) => Alter (ImportTypeWrapper stage) where
  alter ffn fn (ImportTypeWrapper st) =
    ImportTypeWrapper <$> alter ffn fn st

data ImportStatement f a where
  ImportStatement :: Text -> Maybe (ImportList f a) -> a -> ImportStatement f a
  deriving (Generic, Annotated, Alter, Typeable)

data ImportList f a where
  ImportList :: [Identifier f a] -> a -> ImportList f a
  deriving (Generic, Annotated, Alter, Typeable)

-- Top-level declarations.
data FiddleDecl (stage :: Stage) (f :: Type -> Type) a where
  {-
   - An option is a key/value pair.
   - option <ident> <ident>;
   -}
  OptionDecl ::
    Identifier f a ->
    Identifier f a ->
    a ->
    FiddleDecl stage f a
  ImportDecl ::
    ImportType stage f a ->
    a ->
    FiddleDecl stage f a
  UsingDecl ::
    Name f a -> a -> FiddleDecl stage f a
  {- Package Statement. Package Name, Package body -}
  PackageDecl ::
    Name f a ->
    f (PackageBody stage f a) ->
    a ->
    FiddleDecl stage f a
  {- location <identifier> = <expr>. -}
  LocationDecl ::
    Identifier f a ->
    Expression stage f a ->
    a ->
    FiddleDecl stage f a
  {- bits <identifier> : <type> -}
  BitsDecl ::
    Identifier f a ->
    BitType stage f a ->
    a ->
    FiddleDecl stage f a
  {- objtype <identifier> : <type> -}
  ObjTypeDecl ::
    Identifier f a ->
    f (ObjTypeBody stage f a) ->
    a ->
    FiddleDecl stage f a
  {- object <ident> at <expr> : <type> -}
  ObjectDecl ::
    Identifier f a ->
    Expression stage f a ->
    ObjType stage f a ->
    a ->
    FiddleDecl stage f a
  deriving (Generic, Annotated, Typeable)

deriving instance (Alter (ImportType stage)) => Alter (FiddleDecl stage)

data ObjTypeBody (stage :: Stage) (f :: Type -> Type) a where
  ObjTypeBody ::
    BodyType f a ->
    [Directed (ObjTypeDecl stage) f a] ->
    a ->
    ObjTypeBody stage f a
  deriving (Generic, Annotated, Alter, Typeable)

data ObjType stage f a where
  -- { <body> }
  -- Anonymous types are only allowed in stage1. Stage2 should have them be
  -- de-anonymized.
  AnonymousObjType :: f (ObjTypeBody 'Stage1 f a) -> a -> ObjType 'Stage1 f a
  -- <type>[<expr>]
  ArrayObjType :: ObjType stage f a -> Expression stage f a -> a -> ObjType stage f a
  -- <identifier>
  ReferencedObjType :: Name f a -> a -> ObjType stage f a
  deriving (Typeable)

data ObjTypeDecl stage f a where
  {- assert_pos(<expr>) -}
  AssertPosStatement ::
    (StageLessThan stage 3) =>
    Expression stage f a ->
    a ->
    ObjTypeDecl stage f a
  {- reg <ident>(<expr>) : <regtype> -}
  RegisterDecl ::
    Maybe (Modifier f a) ->
    Maybe (Identifier f a) ->
    Expression stage f a ->
    Maybe (RegisterBody stage f a) ->
    a ->
    ObjTypeDecl stage f a
  {- reserved(n); -}
  ReservedDecl :: Expression stage f a -> a -> ObjTypeDecl stage f a
  {- <struct|union> { subfields } <name>; -}
  TypeSubStructure ::
    f (ObjTypeBody stage f a) ->
    Maybe (Identifier f a) ->
    a ->
    ObjTypeDecl stage f a
  deriving (Typeable)

data Modifier f a where
  ModifierKeyword :: ModifierKeyword -> a -> Modifier f a
  deriving (Generic, Annotated, Alter, Typeable)

data ModifierKeyword = Rw | Ro | Wo deriving (Eq, Ord, Show, Read, Typeable)

data DeferredRegisterBody stage f a where
  DeferredRegisterBody ::
    [Directed (RegisterBitsDecl stage) f a] ->
    a ->
    DeferredRegisterBody stage f a
  deriving (Generic, Annotated, Alter, Typeable)

data BodyType (f :: Type -> Type) a where
  Union :: a -> BodyType f a
  Struct :: a -> BodyType f a
  deriving (Generic, Annotated, Alter, Typeable)

data RegisterBody stage f a where
  RegisterBody ::
    BodyType f a ->
    f (DeferredRegisterBody stage f a) ->
    a ->
    RegisterBody stage f a
  deriving (Generic, Annotated, Alter, Typeable)

data RegisterBitsDecl stage f a where
  -- reserved(<expr>)
  ReservedBits :: Expression stage f a -> a -> RegisterBitsDecl stage f a
  -- <modifer> <ident> : <type>
  DefinedBits ::
    Maybe (Modifier f a) ->
    Identifier f a ->
    RegisterBitsTypeRef stage f a ->
    a ->
    RegisterBitsDecl stage f a
  BitsSubStructure ::
    RegisterBody stage f a ->
    Maybe (Identifier f a) ->
    a ->
    RegisterBitsDecl stage f a
  deriving (Generic, Annotated, Alter, Typeable)

data RegisterBitsTypeRef stage f a where
  -- <type>[<expr>]
  RegisterBitsArray ::
    RegisterBitsTypeRef stage f a ->
    Expression stage f a ->
    a ->
    RegisterBitsTypeRef stage f a
  {- Reference to a type. -}
  RegisterBitsReference :: Name f a -> a -> RegisterBitsTypeRef stage f a
  {- enum(<expr>) { <body> }
     Anonymous types are only allowed in stage1.
     Stage2 should de-anonymize these type. -}
  RegisterBitsAnonymousType ::
    AnonymousBitsType Stage1 f a ->
    a ->
    RegisterBitsTypeRef 'Stage1 f a
  {- (<expr>)
   -
   - The expression is just bits ... i.e. an integer.
   -}
  RegisterBitsJustBits ::
    Expression stage f a ->
    a ->
    RegisterBitsTypeRef stage f a

instance Alter (Expression stage) where
  alter ffn fn = \case
    LitNum t a -> LitNum t <$> fn a
    Var i a -> Var <$> alter ffn fn i <*> fn a

instance Annotated (Expression stage) where
  annot = \case
    LitNum _ a -> a
    Var _ a -> a

instance Alter (ObjTypeDecl stage) where
  alter ffn fn = \case
    (AssertPosStatement expr a) ->
      AssertPosStatement <$> alter ffn fn expr <*> fn a
    (RegisterDecl mMod mIdent expr mBody a) ->
      RegisterDecl
        <$> mapM (alter ffn fn) mMod
        <*> mapM (alter ffn fn) mIdent
        <*> alter ffn fn expr
        <*> mapM (alter ffn fn) mBody
        <*> fn a
    (TypeSubStructure mBody mIdent a) ->
      TypeSubStructure
        <$> (ffn =<< mapM (alter ffn fn) mBody)
        <*> mapM (alter ffn fn) mIdent
        <*> fn a
    (ReservedDecl expr a) ->
      ReservedDecl
        <$> alter ffn fn expr
        <*> fn a

instance Annotated (ObjTypeDecl stage) where
  annot = \case
    (AssertPosStatement _ a) -> a
    (RegisterDecl _ _ _ _ a) -> a

instance Alter (ObjType stage) where
  alter ffn fn = \case
    (AnonymousObjType b a) ->
      AnonymousObjType <$> (ffn =<< mapM (alter ffn fn) b) <*> fn a
    (ArrayObjType t e a) ->
      ArrayObjType <$> alter ffn fn t <*> alter ffn fn e <*> fn a
    (ReferencedObjType i a) ->
      ReferencedObjType <$> alter ffn fn i <*> fn a

instance Annotated (ObjType stage) where
  annot = \case
    (AnonymousObjType _ a) -> a
    (ArrayObjType _ _ a) -> a
    (ReferencedObjType _ a) -> a

instance Alter (RegisterBitsTypeRef stage) where
  alter ffn fn = \case
    (RegisterBitsArray ref exp a) ->
      RegisterBitsArray <$> alter ffn fn ref <*> alter ffn fn exp <*> fn a
    (RegisterBitsReference i a) ->
      RegisterBitsReference <$> alter ffn fn i <*> fn a
    (RegisterBitsAnonymousType t a) ->
      RegisterBitsAnonymousType <$> alter ffn fn t <*> fn a
    (RegisterBitsJustBits e a) ->
      RegisterBitsJustBits <$> alter ffn fn e <*> fn a

instance Annotated (RegisterBitsTypeRef stage) where
  annot = \case
    (RegisterBitsArray _ _ a) -> a
    (RegisterBitsReference _ a) -> a
    (RegisterBitsAnonymousType _ a) -> a
    (RegisterBitsJustBits _ a) -> a

data AnonymousBitsType stage f a where
  -- enum(<expr>) { <body> }
  AnonymousEnumBody ::
    Expression stage f a ->
    f (EnumBody stage f a) ->
    a ->
    AnonymousBitsType stage f a
  deriving (Generic, Annotated, Alter, Typeable)

data BitType (stage :: Stage) (f :: Type -> Type) a where
  -- enum(<expr>) { <body> }
  EnumBitType ::
    Expression stage f a ->
    f (EnumBody stage f a) ->
    a ->
    BitType stage f a
  -- (<expr>)
  RawBits :: Expression stage f a -> a -> BitType stage f a
  deriving (Generic, Annotated, Alter, Typeable)

data EnumBody (stage :: Stage) (f :: Type -> Type) a where
  -- <decl>,
  EnumBody :: [Directed (EnumConstantDecl stage) f a] -> a -> EnumBody stage f a
  deriving (Generic, Annotated, Alter, Typeable)

data EnumConstantDecl stage f a where
  -- <ident> = <expr>
  EnumConstantDecl ::
    Identifier f a ->
    Expression stage f a ->
    a ->
    EnumConstantDecl stage f a
  -- reserved = <expr>
  EnumConstantReserved ::
    Expression stage f a ->
    a ->
    EnumConstantDecl stage f a
  deriving (Generic, Annotated, Alter, Typeable)

data PackageBody (stage :: Stage) (f :: Type -> Type) a where
  {- The body of a package -}
  PackageBody :: [Directed (FiddleDecl stage) f a] -> a -> PackageBody stage f a
  deriving (Generic, Annotated, Typeable)

deriving instance (Alter (ImportType stage)) => Alter (PackageBody stage)

-- instance Alter (Modifier stage) where
--   alter _ fn (ModifierKeyword m a) = ModifierKeyword m (fn a)
--
-- instance Alter (Identifier) where
--   alter _ fn (Identifier i a) = Identifier i $ fn a
--
-- instance Alter (Expression stage) where
--   alter ffn fn = \case
--     (LitNum t a) -> LitNum t $ fn a
--     (RealNum t a) -> RealNum t $ fn a
--     (Var i a) -> Var (alter ffn fn i) $ fn a

proxyOf :: t f a -> Proxy t
proxyOf _ = Proxy

class Annotated (t :: SynTreeKind) where
  annot :: t f a -> a
  default annot :: (Generic (t f a), GAnnot a (Rep (t f a))) => t f a -> a
  annot t = gannot (from t)

class GAnnot a r where
  gannot :: r x -> a

instance GAnnot a (Rec0 a) where
  gannot = unK1

instance (GAnnot a r) => GAnnot a (l :*: r) where
  gannot (_ :*: r) = gannot r

instance (GAnnot a r, GAnnot a l) => GAnnot a (l :+: r) where
  gannot (R1 r) = gannot r
  gannot (L1 l) = gannot l

instance (GAnnot a r) => GAnnot a (M1 i c r) where
  gannot (M1 a) = gannot a

class Alter (t :: SynTreeKind) where
  alter ::
    forall f1 f2 a1 a2 m.
    (Monad m, Traversable f1) =>
    (forall z. f1 z -> m (f2 z)) ->
    (a1 -> m a2) ->
    t f1 a1 ->
    m (t f2 a2)
  default alter ::
    forall f1 f2 a1 a2 m.
    ( Generic (t f1 a1),
      Generic (t f2 a2),
      Traversable f1,
      GAlter t f1 f2 a1 a2 (Rep (t f1 a1)) (Rep (t f2 a2)),
      Monad m
    ) =>
    ( forall z.
      f1 z ->
      m (f2 z)
    ) ->
    (a1 -> m a2) ->
    t f1 a1 ->
    m (t f2 a2)
  alter ffn fn t = to <$> galter (proxyOf t) ffn fn (from t)

instance (Alter t, Traversable f) => Functor (t f) where
  fmap f t = runIdentity (alter return (return . f) t)

class GAlter t f1 f2 a1 a2 r1 r2 where
  galter ::
    forall proxy x m.
    (Monad m, Traversable f1) =>
    proxy t ->
    (forall z. f1 z -> m (f2 z)) ->
    (a1 -> m a2) ->
    r1 x ->
    m (r2 x)

{- Altering a record with type a1 will apply the mapping function and produce a
    record with type a2 -}
instance GAlter t f1 f2 a1 a2 (Rec0 a1) (Rec0 a2) where
  galter _ _ fn k1 = K1 <$> fn (unK1 k1)

{- Base-Case. Altering unrelated leaf types will do nothing. -}
instance GAlter t f1 f2 a1 a2 (Rec0 u1) (Rec0 u1) where
  galter _ _ _ = return

{- Recursive case. Call alter on sub-structures. -}
instance
  (Alter u) =>
  GAlter t f1 f2 a1 a2 (Rec0 (u f1 a1)) (Rec0 (u f2 a2))
  where
  galter _ ffn fn k1 =
    K1 <$> alter ffn fn (unK1 k1)

{- Recursive case. Called when there are list of substructures that need to be
    recused. -}
instance
  (Alter u, Traversable l) =>
  GAlter t f1 f2 a1 a2 (Rec0 (l (u f1 a1))) (Rec0 (l (u f2 a2)))
  where
  galter _ ffn fn k1 =
    K1 <$> mapM (alter ffn fn) (unK1 k1)

-- instance GAlter t f1 f2 a1 a2 (Rec0 (f1 z)) (Rec0 (f2 z)) where
--   galter _ ffn _ k1 = K1 <$> ffn (unK1 k1)

{- Generic altering. Descends into the function and alters whatever is inside
    the functor and then transforms the functor using the ffn function. -}
instance
  (Traversable f1, Alter u) =>
  GAlter t f1 f2 a1 a2 (Rec0 (f1 (u f1 a1))) (Rec0 (f2 (u f2 a2)))
  where
  galter proxy ffn fn k1 = do
    newK <- mapM (alter ffn fn) (unK1 k1)
    K1 <$> ffn newK

instance
  ( GAlter t f1 f2 a1 a2 l1 l2,
    GAlter t f1 f2 a1 a2 r1 r2
  ) =>
  GAlter t f1 f2 a1 a2 (l1 :*: r1) (l2 :*: r2)
  where
  galter proxy ffn fn (a :*: b) = do
    a' <- galter proxy ffn fn a
    b' <- galter proxy ffn fn b
    return (a' :*: b')

instance
  ( GAlter t f1 f2 a1 a2 l1 l2,
    GAlter t f1 f2 a1 a2 r1 r2
  ) =>
  GAlter t f1 f2 a1 a2 (l1 :+: r1) (l2 :+: r2)
  where
  galter proxy ffn fn (R1 r) = R1 <$> galter proxy ffn fn r
  galter proxy ffn fn (L1 l) = L1 <$> galter proxy ffn fn l

instance
  (GAlter t f1 f2 a1 a2 r1 r2) =>
  GAlter t f1 f2 a1 a2 (M1 i c r1) (M1 i c r2)
  where
  galter proxy ffn fn (M1 a) = M1 <$> galter proxy ffn fn a

type family StageNumber (s :: Stage) :: Natural where
  StageNumber Stage1 = 1
  StageNumber Stage2 = 2
  StageNumber Stage3 = 3

{--}
squeeze :: (Alter t, Traversable f, Monad f) => t f a -> f (t Identity a)
squeeze = alter (fmap Identity) return