add datatype generation DSL stuff
[clean-tests.git] / datatype / MkCons.hs
diff --git a/datatype/MkCons.hs b/datatype/MkCons.hs
new file mode 100644 (file)
index 0000000..57b14c6
--- /dev/null
@@ -0,0 +1,173 @@
+{-# LANGUAGE TemplateHaskell #-}
+module MkCons where
+
+import Language.Haskell.TH.Syntax
+import Language.Haskell.TH
+import Data.Char
+import Control.Monad
+
+className :: Name -> Name
+className = mkName . (++"'") . stringName
+constructorName :: Name -> Name
+constructorName = mkName . map toLower . stringName
+selectorName :: Name -> Name
+selectorName = mkName . map toLower . (++"'") . stringName
+stringName :: Name -> String
+stringName (Name occ _) = occString occ
+
+numberedArgs :: [a] -> [String]
+numberedArgs = zipWith (\i _->"f" ++ show i) [0 :: Int ..]
+
+mkConsClass :: Name -> DecsQ
+mkConsClass typename = reify typename >>= \info->case info of
+    TyConI dec
+        -> case dec of
+        DataD _ _ tyvars _ constructors _
+            -> sequence 
+                [ mkDerivation tyvars
+                , mkConstructorClasses tyvars constructors
+                , mkPrinterInstances tyvars constructors
+                , mkCompilerInstances tyvars constructors
+                ]
+        _
+            -> fail "mkConsClass only supports data types"
+    _
+        -> fail "mkConsClass only supports types"
+  where
+    mkDerivation :: [TyVarBndr] -> DecQ
+    mkDerivation tyvars = let names = (map (VarT . getNameTyVarBndr) tyvars) in pure $
+        InstanceD Nothing
+            [ConT (mkName "Serialise") `AppT` t | t <- names]
+            (ConT (mkName "Serialise") `AppT` foldl AppT (ConT typename) names)
+            []
+
+    mkConstructorClasses :: [TyVarBndr] -> [Con] -> DecQ
+    mkConstructorClasses tyvars constructors = do
+        cclasses <- mapM mkConstructorClassMember constructors
+        sclasses <- concat <$> mapM mkSelectorClassMember constructors
+        pure $ ClassD [] (className typename) [PlainTV view] [] (cclasses ++ sclasses)
+      where
+        view = mkName "m"
+
+        mkConstructorClassMember :: Con -> DecQ
+        mkConstructorClassMember (NormalC consname fs)
+            = mkConstructorClassMemberForName consname [t | (_, t)<-fs]
+        mkConstructorClassMember (RecC consname fs)
+            = mkConstructorClassMemberForName consname [t | (_, _, t)<-fs]
+        mkConstructorClassMember t
+            = fail $ "mkConsClass not supported for types such as: " ++ show t
+
+        mkConstructorClassMemberForName :: Name -> [Type] -> DecQ
+        mkConstructorClassMemberForName consname fs
+            = pure $ SigD (constructorName consname)
+                $ foldr (AppT . AppT ArrowT) resultT
+                $ map (AppT $ VarT view) fs
+
+        mkSelectorClassMember :: Con -> DecsQ
+        mkSelectorClassMember (NormalC _ fs)
+            = mapM (uncurry mkSelectorClassMemberForField)
+            $ zipWith (\(_, t) i->(mkName $ map toLower (stringName typename) ++ "f" ++ show i, t)) fs [0 :: Int ..]
+        mkSelectorClassMember (RecC _ fs)
+            = mapM (uncurry mkSelectorClassMemberForField) [(n, t)|(n, _, t)<-fs]
+        mkSelectorClassMember t
+            = fail $ "mkConsClass not supported for types such as: " ++ show t
+
+        mkSelectorClassMemberForField :: Name -> Type -> DecQ
+        mkSelectorClassMemberForField n t = pure
+            $ SigD (className n)
+            $ ForallT [] [ConT (mkName "Serialise") `AppT` VarT (getNameTyVarBndr ty) | ty <- tyvars]
+            $ ArrowT `AppT` resultT `AppT` (VarT view `AppT` t)
+
+        resultT = VarT view `AppT` (foldl AppT (ConT typename) $ map (VarT . getNameTyVarBndr) $ tyvars)
+
+    mkPrinterInstances :: [TyVarBndr] -> [Con] -> DecQ
+    mkPrinterInstances _ constructors
+        =   InstanceD Nothing [] (AppT (ConT $ className typename) (ConT $ mkName "Printer")) . concat
+        <$> mapM mkPrinterInstance constructors
+      where
+        mkPrinterInstance :: Con -> DecsQ
+        mkPrinterInstance (NormalC name fs)
+            | null fs = pure [FunD (constructorName name)
+                    [Clause [] (NormalB $ VarE (mkName "printLit") `AppE` LitE (StringL $ stringName name)) [] ]]
+            | otherwise =
+                let args = map mkName $ numberedArgs fs
+                in (:) <$> pure (FunD (constructorName name)
+                    [Clause
+                        (map VarP args)
+                        (NormalB $
+                            (AppE (VarE (mkName "printCons") `AppE` LitE (StringL $ stringName name))
+                                (foldl (\x y->x `pc` pl " " `pc` y) (pl "") $ map VarE args)
+                            )
+                        )
+                        []
+                    ])
+                <*> mapM mkPrinterSelector
+                    (zipWith (\_ i->map toLower (stringName typename) ++ "f" ++ show i) fs [0 :: Int ..])
+        mkPrinterInstance (RecC name fs)
+            = let args = map mkName $ numberedArgs fs
+            in (:) <$> pure (FunD (constructorName name)
+                [Clause
+                    (map VarP args)
+                    (NormalB $
+                        (AppE (VarE (mkName "printRec") `AppE` LitE (StringL $ stringName typename))
+                            (foldl1 (\x y->x `pc` pl ", " `pc` y)
+                                $ zipWith (\x ((Name occ'' _), _, _)->pl (occString occ'') `pc` pl "=" `pc` VarE x) args fs)
+                        )
+                    )
+                    []
+                ])
+            <*> mapM mkPrinterSelector [occString occ | (Name occ _, _, _)<-fs]
+        mkPrinterInstance t = fail $ "mkConsClass not supported for types such as: " ++ show t
+
+        mkPrinterSelector :: String -> Q Dec
+        mkPrinterSelector n' = do
+            body <- [|\d->d >> printLit $(pure $ LitE $ StringL ('.':n'))|]
+            pure $ FunD (selectorName $ mkName n') [Clause [] (NormalB body) []]
+
+    mkCompilerInstances :: [TyVarBndr] -> [Con] -> DecQ
+    mkCompilerInstances _ constructors
+        =   InstanceD Nothing [] (AppT (ConT $ className typename) (ConT $ mkName "Compiler")) . concat
+        <$> mapM (uncurry mkCompilerInstance) (zip constructors [0..])
+      where
+        mkCompilerInstance :: Con -> Int -> DecsQ
+        mkCompilerInstance (NormalC name fs) consnum = (:)
+            <$> mkCompilerInstanceForName name consnum (numberedArgs fs)
+            <*> mapM (uncurry mkCompilerSelector) (zip [0..] [map toLower (stringName typename) ++ f | f<-numberedArgs fs])
+        mkCompilerInstance (RecC name fs) consnum = (:)
+            <$> mkCompilerInstanceForName name consnum [occString occ | (Name occ _, _, _) <- fs]
+            <*> mapM (uncurry mkCompilerSelector) (zip [0..] [occString occ | (Name occ _, _, _) <- fs])
+        mkCompilerInstance t _ = fail $ "mkConsClass not supported for types such as: " ++ show t
+
+        mkCompilerInstanceForName :: Name -> Int -> [String] -> DecQ
+        mkCompilerInstanceForName name consnum fs =
+            let args = map mkName $ numberedArgs fs
+            in do
+                body <- [|instr [Push $(pure $ LitE $ IntegerL $ toInteger $ consnum)] *> $(mkBody $ map VarE args) <* instr [Sth $ $(pure $ LitE $ IntegerL $ toInteger $ length fs) + 1] |]
+                pure $ FunD (constructorName name)
+                    [Clause (map VarP args) (NormalB body) [] ]
+          where
+            mkBody :: [Exp] -> Q Exp
+            mkBody [] = pure $ VarE (mkName "pure") `AppE` ConE name
+            mkBody (a:as) = foldM (\x y->pure $ ifx "<*>" x y) (ifx "<$>" (ConE name) a) as
+
+        mkCompilerSelector :: Int -> String -> DecQ
+        mkCompilerSelector idx n' = do
+            body <- [|\d->d >> instr [Ldh $(pure $ LitE (IntegerL (toInteger idx)))] |]
+            pure $ FunD (selectorName $ mkName n')
+                [Clause [] (NormalB body) [] ]
+
+instrE :: Exp -> Exp
+instrE e = VarE (mkName "instr") `AppE` ListE [e]
+
+ifx :: String -> Exp -> Exp -> Exp
+ifx op a b = InfixE (Just a) (VarE $ mkName op) (Just b)
+
+pc :: Exp -> Exp -> Exp
+pc l r = VarE (mkName ">>>") `AppE` l `AppE` r
+
+pl :: String -> Exp
+pl s = VarE (mkName "printLit") `AppE` LitE (StringL s)
+
+getNameTyVarBndr :: TyVarBndr -> Name
+getNameTyVarBndr (PlainTV name) = name
+getNameTyVarBndr (KindedTV name _) = name