1 files changed, 259 insertions, 51 deletions

diff --git a/src/HSVIS/AST.hs b/src/HSVIS/AST.hs
index f606d22..7b05e19 100644
--- a/src/HSVIS/AST.hs
+++ b/src/HSVIS/AST.hs
@@ -1,75 +1,162 @@
+{-# LANGUAGE ConstraintKinds #-}
 {-# LANGUAGE DataKinds #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE TypeApplications #-}
+{-# LANGUAGE LambdaCase #-}
 module HSVIS.AST where
 
+import Data.Bifunctor (bimap, second)
+import qualified Data.Kind as DK
 import Data.List.NonEmpty (NonEmpty)
+import Data.Proxy
 
 import HSVIS.Pretty
 
 
+-- | An AST type has a single stage argument.
+type ASTType = DK.Type -> DK.Type
+
+-- | Trees that Grow extension type family
+type family X (f :: ASTType) s
+
+-- | Additional constructors
+data family E (f :: ASTType) s
+
+class HasExt f where
+  type HasXField f :: Bool
+  type HasECon f :: Bool
+  type ContainsData f :: [ASTType]
+  extOf :: f s -> X f s
+  extMap :: Proxy f -> Proxy s1 -> Proxy s2
+         -> XMapperTelescope (TransitiveContainsData '[f]) s1 s2
+              (EMapperTelescope (TransitiveContainsData '[f]) s1 s2
+                (f s1 -> f s2))
+
+type family If c a b where
+  If 'True a b = a
+  If 'False a b = b
+
+type family ASTShowPrereqs f s :: DK.Constraint where
+  ASTShowPrereqs f s =
+    (If (HasXField f) (Show (X f s)) (() :: DK.Constraint)
+    ,If (HasECon f) (Show (E f s)) (() :: DK.Constraint)
+    ,ASTShowPrereqsC (ContainsData f) s)
+
+type family ASTShowPrereqsC l s :: DK.Constraint where
+  ASTShowPrereqsC '[] s = ()
+  ASTShowPrereqsC (f ': fs) s = (Show (f s), ASTShowPrereqsC fs s)
+
+type family a && b where
+  'False && b = 'False
+  'True && b = b
+
+type family In x l where
+  In x '[] = 'False
+  In x (x ': ys) = 'True
+  In x (_ ': ys) = In x ys
+
+type family Sublist l1 l2 where
+  Sublist '[] l = 'True
+  Sublist (x ': xs) l = In x l && Sublist xs l
+
+type family Union l1 l2 where
+  Union '[] l = l
+  Union (x ': xs) l = If (In x l) (Union xs l) (x ': Union xs l)
+
+type family MapContainsData l where
+  MapContainsData '[] = '[]
+  MapContainsData (f ': fs) = Union (ContainsData f) (MapContainsData fs)
+
+type family TransitiveContainsData fs where
+  TransitiveContainsData fs = TransitiveContainsData' fs (Sublist (Union fs (MapContainsData fs)) fs)
+
+type family TransitiveContainsData' fs isfull where
+  TransitiveContainsData' fs 'True = fs
+  TransitiveContainsData' fs 'False = TransitiveContainsData (Union fs (MapContainsData fs))
+
+type family XMapperTelescope fs s1 s2 a where
+  XMapperTelescope '[] s1 s2 a = a
+  XMapperTelescope (f ': fs) s1 s2 a =
+    If (HasXField f)
+       ((X f s1 -> X f s2) -> XMapperTelescope fs s1 s2 a)
+       (XMapperTelescope fs s1 s2 a)
+
+type family EMapperTelescope fs s1 s2 a where
+  EMapperTelescope '[] s1 s2 a = a
+  EMapperTelescope (f ': fs) s1 s2 a =
+    If (HasECon f)
+       ((X f s1 -> E f s1 -> f s2) -> EMapperTelescope fs s1 s2 a)
+       (EMapperTelescope fs s1 s2 a)
+
+
 newtype Name = Name String
   deriving (Show, Eq)
 
-data Program t = Program [DataDef] [FunDef t]
-  deriving (Show)
+data Program s = Program [DataDef s] [FunDef s]
+deriving instance (Show (DataDef s), Show (FunDef s)) => Show (Program s)
 
-data DataDef = DataDef Name [Name] [(Name, [Type])]
-  deriving (Show)
+data DataDef s = DataDef (X DataDef s) Name [Name] [(Name, [Type s])]
+deriving instance ASTShowPrereqs DataDef s => Show (DataDef s)
 
-data FunDef t = FunDef Name (Maybe Type) (NonEmpty (FunEq t))
-  deriving (Show)
+data FunDef s = FunDef (X FunDef s) Name (Maybe (Type s)) (NonEmpty (FunEq s))
+deriving instance ASTShowPrereqs FunDef s => Show (FunDef s)
 
-data FunEq t = FunEq t Name [Pattern t] (RHS t)
-  deriving (Show)
+data FunEq s = FunEq (X FunEq s) Name [Pattern s] (RHS s)
+deriving instance ASTShowPrereqs FunEq s => Show (FunEq s)
 
-data TypeStage = TSTC | TSNormal
-  deriving (Show)
+data Kind s
+  = KType (X Kind s)
+  | KFun (X Kind s) (Kind s) (Kind s)
+deriving instance ASTShowPrereqs Kind s => Show (Kind s)
 
-data GType (stage :: TypeStage) where
-  TApp :: GType s -> [GType s] -> GType s
-  TTup :: [GType s] -> GType s
-  TList :: GType s -> GType s
-  TFun :: GType s -> GType s -> GType s
-  TCon :: Name -> GType s
-  TVar :: Name -> GType s
-
-  -- Type constructor used only while type checking
-  TUniVar :: Int -> GType 'TSTC
-deriving instance Show (GType stage)
-
-type TCType = GType 'TSTC
-type Type = GType 'TSNormal
-
-data Pattern t
-  = PWildcard t
-  | PVar t Name
-  | PAs t Name (Pattern t)
-  | PCon t Name [Pattern t]
-  | POp t (Pattern t) Operator (Pattern t)
-  | PList t [Pattern t]
-  | PTup t [Pattern t]
-  deriving (Show)
+data Type s
+  = TApp (X Type s) (Type s) [Type s]
+  | TTup (X Type s) [Type s]
+  | TList (X Type s) (Type s)
+  | TFun (X Type s) (Type s) (Type s)
+  | TCon (X Type s) Name
+  | TVar (X Type s) Name
 
-data RHS t
-  = Guarded [(Expr t, Expr t)]  -- currently not parsed
-  | Plain (Expr t)
-  deriving (Show)
+  -- extension point
+  | TExt (X Type s) !(E Type s)
+deriving instance ASTShowPrereqs Type s => Show (Type s)
 
-data Expr t
-  = ELit t Literal
-  | EVar t Name
-  | ECon t Name
-  | EList t [Expr t]
-  | ETup t [Expr t]
-  | EApp t (Expr t) [Expr t]
-  | EOp t (Expr t) Operator (Expr t)
-  | EIf t (Expr t) (Expr t) (Expr t)
-  | ECase t (Expr t) [(Pattern t, RHS t)]
-  | ELet t [FunDef t] (Expr t)
-  | EParseError
-  deriving (Show)
+data Pattern s
+  = PWildcard (X Pattern s)
+  | PVar      (X Pattern s) Name
+  | PAs       (X Pattern s) Name (Pattern s)
+  | PCon      (X Pattern s) Name [Pattern s]
+  | POp       (X Pattern s) (Pattern s) Operator (Pattern s)
+  | PList     (X Pattern s) [Pattern s]
+  | PTup      (X Pattern s) [Pattern s]
+deriving instance ASTShowPrereqs Pattern s => Show (Pattern s)
+
+data RHS s
+  = Guarded (X RHS s) [(Expr s, Expr s)]  -- currently not parsed
+  | Plain   (X RHS s) (Expr s)
+deriving instance ASTShowPrereqs RHS s => Show (RHS s)
+
+data Expr s
+  = ELit  (X Expr s) Literal
+  | EVar  (X Expr s) Name
+  | ECon  (X Expr s) Name
+  | EList (X Expr s) [Expr s]
+  | ETup  (X Expr s) [Expr s]
+  | EApp  (X Expr s) (Expr s) [Expr s]
+  | EOp   (X Expr s) (Expr s) Operator (Expr s)
+  | EIf   (X Expr s) (Expr s) (Expr s) (Expr s)
+  | ECase (X Expr s) (Expr s) [(Pattern s, RHS s)]
+  | ELet  (X Expr s) [FunDef s] (Expr s)
+  | EError (X Expr s)
+deriving instance ASTShowPrereqs Expr s => Show (Expr s)
 
 data Literal = LInt Integer | LFloat Rational | LChar Char | LString String
   deriving (Show)
@@ -80,3 +167,124 @@ data Operator = OAdd | OSub | OMul | ODiv | OMod | OEqu | OPow
 
 instance Pretty Name where
   prettysPrec _ (Name n) = showString ("\"" ++ n ++ "\"")
+
+instance HasExt DataDef where
+  type HasXField DataDef = 'True
+  type HasECon DataDef = 'False
+  type ContainsData DataDef = '[Type]
+  extOf (DataDef x _ _ _) = x
+  extMap _ ps1 ps2 fdd fty gty (DataDef x n ps cs) =
+    DataDef (fdd x) n ps (map (second (map (extMap (Proxy @Type) ps1 ps2 fty gty))) cs)
+
+instance HasExt FunDef where
+  type HasXField FunDef = 'True
+  type HasECon FunDef = 'False
+  type ContainsData FunDef = '[Type, FunEq]
+  extOf (FunDef x _ _ _) = x
+  extMap _ ps1 ps2 ffd ffe frh fty fpa fex gty (FunDef x n msig es) =
+    FunDef (ffd x) n (extMap (Proxy @Type) ps1 ps2 fty gty <$> msig)
+           (fmap (extMap (Proxy @FunEq) ps1 ps2 frh ffd fpa fex fty ffe gty) es)
+
+instance HasExt FunEq where
+  type HasXField FunEq = 'True
+  type HasECon FunEq = 'False
+  type ContainsData FunEq = '[Pattern, RHS]
+  extOf (FunEq x _ _ _) = x
+  extMap _ ps1 ps2 frh ffd fpa fex fty ffe gty (FunEq x n ps rhs) =
+    FunEq (ffe x) n (map (extMap (Proxy @Pattern) ps1 ps2 fpa) ps)
+          (extMap (Proxy @RHS) ps1 ps2 frh ffd fpa fex fty ffe gty rhs)
+
+instance HasExt Kind where
+  type HasXField Kind = 'True
+  type HasECon Kind = 'False
+  type ContainsData Kind = '[Kind]
+  extOf (KType x) = x
+  extOf (KFun x _ _) = x
+  extMap _ _   _   f (KType x) = KType (f x)
+  extMap p ps1 ps2 f (KFun x a b) = KFun (f x) (extMap p ps1 ps2 f a) (extMap p ps1 ps2 f b)
+
+instance HasExt Type where
+  type HasXField Type = 'True
+  type HasECon Type = 'True
+  type ContainsData Type = '[Type]
+  extOf (TApp x _ _) = x
+  extOf (TTup x _) = x
+  extOf (TList x _) = x
+  extOf (TFun x _ _) = x
+  extOf (TCon x _) = x
+  extOf (TVar x _) = x
+  extOf (TExt x _) = x
+
+  extMap p ps1 ps2 f g (TApp x a b) = TApp (f x) (extMap p ps1 ps2 f g a) (map (extMap p ps1 ps2 f g) b)
+  extMap p ps1 ps2 f g (TTup x as) = TTup (f x) (map (extMap p ps1 ps2 f g) as)
+  extMap p ps1 ps2 f g (TList x a) = TList (f x) (extMap p ps1 ps2 f g a)
+  extMap p ps1 ps2 f g (TFun x a b) = TFun (f x) (extMap p ps1 ps2 f g a) (extMap p ps1 ps2 f g b)
+  extMap _ _   _   f _ (TCon x n) = TCon (f x) n
+  extMap _ _   _   f _ (TVar x n) = TVar (f x) n
+  extMap _ _   _   _ g (TExt x e) = g x e
+
+instance HasExt Pattern where
+  type HasXField Pattern = 'True
+  type HasECon Pattern = 'False
+  type ContainsData Pattern = '[Pattern]
+  extOf (PWildcard x) = x
+  extOf (PVar x _) = x
+  extOf (PAs x _ _) = x
+  extOf (PCon x _ _) = x
+  extOf (POp x _ _ _) = x
+  extOf (PList x _) = x
+  extOf (PTup x _) = x
+
+  extMap _ _   _   f (PWildcard x) = PWildcard (f x)
+  extMap _ _   _   f (PVar x n) = PVar (f x) n
+  extMap p ps1 ps2 f (PAs x n a) = PAs (f x) n (extMap p ps1 ps2 f a)
+  extMap p ps1 ps2 f (PCon x n as) = PCon (f x) n (map (extMap p ps1 ps2 f) as)
+  extMap p ps1 ps2 f (POp x a n b) = POp (f x) (extMap p ps1 ps2 f a) n (extMap p ps1 ps2 f b)
+  extMap p ps1 ps2 f (PList x as) = PList (f x) (map (extMap p ps1 ps2 f) as)
+  extMap p ps1 ps2 f (PTup x as) = PTup (f x) (map (extMap p ps1 ps2 f) as)
+
+instance HasExt RHS where
+  type HasXField RHS = 'True
+  type HasECon RHS = 'False
+  type ContainsData RHS = '[Expr]
+  extOf (Guarded x _) = x
+  extOf (Plain x _) = x
+
+  extMap _ ps1 ps2 frh ffd fpa fex fty ffe gty (Guarded x l) =
+    Guarded (frh x) (map (bimap (extMap (Proxy @Expr) ps1 ps2 frh ffd fpa fex fty ffe gty)
+                                (extMap (Proxy @Expr) ps1 ps2 frh ffd fpa fex fty ffe gty))
+                         l)
+  extMap _ ps1 ps2 frh ffd fpa fex fty ffe gty (Plain x e) =
+    Plain (frh x) (extMap (Proxy @Expr) ps1 ps2 frh ffd fpa fex fty ffe gty e)
+
+instance HasExt Expr where
+  type HasXField Expr = 'True
+  type HasECon Expr = 'False
+  type ContainsData Expr = '[Expr, Pattern, RHS, FunDef]
+  extOf (ELit x _) = x
+  extOf (EVar x _) = x
+  extOf (ECon x _) = x
+  extOf (EList x _) = x
+  extOf (ETup x _) = x
+  extOf (EApp x _ _) = x
+  extOf (EOp x _ _ _) = x
+  extOf (EIf x _ _ _) = x
+  extOf (ECase x _ _) = x
+  extOf (ELet x _ _) = x
+  extOf (EError x) = x
+
+  extMap p ps1 ps2 frh ffd fpa fex fty ffe gty = \case
+    ELit x l -> ELit (fex x) l
+    EVar x n -> EVar (fex x) n
+    ECon x n -> ECon (fex x) n
+    EList x es -> EList (fex x) (map (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty) es)
+    ETup x es -> ETup (fex x) (map (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty) es)
+    EApp x a es -> EApp (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a) (map (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty) es)
+    EOp x a n b -> EOp (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a) n (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty b)
+    EIf x a b c -> EIf (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty b) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty c)
+    ECase x a l -> ECase (fex x) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty a)
+                         (map (bimap (extMap (Proxy @Pattern) ps1 ps2 fpa)
+                                     (extMap (Proxy @RHS) ps1 ps2 frh ffd fpa fex fty ffe gty))
+                              l)
+    ELet x as b -> ELet (fex x) (map (extMap (Proxy @FunDef) ps1 ps2 ffd ffe frh fty fpa fex gty) as) (extMap p ps1 ps2 frh ffd fpa fex fty ffe gty b)
+    EError x -> EError (fex x)