Files
rlp/tst/CoreSyntax.hs
2023-12-29 22:29:04 -07:00

304 lines
10 KiB
Haskell

{-# LANGUAGE OverloadedStrings, LambdaCase, GeneralisedNewtypeDeriving #-}
module CoreSyntax
( ProgramSrc(..)
, congruentSrc
, unparseCoreProg
)
where
----------------------------------------------------------------------------------
import Core.Syntax
import Compiler.JustRun (justParseSrc)
import Control.Arrow ((>>>), (&&&))
import Control.Monad
import Data.List (intersperse)
import Data.Coerce (coerce)
import Data.Text (Text)
import Data.Text qualified as T
import Data.HashMap.Strict qualified as H
import Test.QuickCheck
import Text.PrettyPrint hiding ((<>))
import Data.Functor ((<&>))
import Data.Function ((&), on)
import Data.String (IsString(..))
import Lens.Micro.Platform
import Lens.Micro.Platform.Internal (IsText(..))
----------------------------------------------------------------------------------
newtype ProgramSrc = ProgramSrc Text
deriving (Show, Read, Eq, Semigroup, Monoid, IsString)
instance Arbitrary ProgramSrc where
arbitrary = sized genProg where
genProg :: Int -> Gen ProgramSrc
genProg n = do
-- in generating a program, we create a random list of sc names and
-- assign them type signatures and definitions in random order.
ns <- replicateM n genName
-- generate a typesig and def for each name
ns & each %~ (genTySig &&& genScDef)
-- [(typesig, scdef)] -> [typesigs and scdefs]
& uncurry (++) . unzip
-- [Gen Text] -> Gen [Text]
& sequenceA
-- shuffle order of tysigs and scdefs
>>= shuffle
-- terminate each tysig and scdef with a semicolon with a blank
-- line for legibility
<&> intersperse ";\n\n"
-- mconcat into a single body of text
<&> mconcat
-- she's done! put a bow on her! :D
<&> ProgramSrc
genTySig :: Name -> Gen Text
genTySig n = conseq [pure n, ws, pure "::", ws, genTy]
genScDef :: Name -> Gen Text
genScDef n = conseq [pure n, ws, pure "=", ws, genExpr]
genExpr :: Gen Text
genExpr = gen 4 0 where
gen 0 _ = oneof
[ genVar
, genLit
]
gen n p = oneof
[ gen 0 p
, wrapParens <$> gen n' 0
, genApp n p
, genLet n p
-- , genLam n p
-- , genCase n p
]
where n' = next n
genVar = oneof
[ genName
, genCon
, wrapParens <$> genSymName
, wrapParens <$> genSymCon
]
genCase n p = conseq [ pure "case", ws1, gen n' 0, ws1, pure "of"
, pure "{", alts, pure "}"
]
<&> pprec 0 p
where
n' = next n
alts = chooseSize (1,6) (listOf1 alt)
<&> intersperse ";"
<&> mconcat
alt = conseq [ tag, ws, pure "->", ws1, gen n' 0 ]
tag = T.pack . show <$> chooseInt (0,maxBound)
genLit = T.pack . show <$> chooseInt (0,maxBound)
genApp n p = chooseSize (2,10) (listOf1 (gen n' 1))
<&> pprec 0 p . mconcat . intersperse " "
where
n' = next n
genLet n p = conseq [ letw, ws, pure "{", ws, binds
, ws, pure "}", ws, pure "in"
, ws1, gen n' 0
]
where
letw = arbitrary <&> \case
Rec -> "letrec"
NonRec -> "let"
binds = chooseSize (1,6) (listOf1 bind)
<&> intersperse ";"
<&> mconcat
bind = conseq [var, ws, pure "=", ws, gen n' 0]
var = oneof [genName, wrapParens <$> genSymName]
n' = next n
genLam n p = conseq [l, ws, bs, ws, pure "->", ws, gen n' 0]
<&> pprec 0 p
where
-- whitespace because reserved op shenanigans :3
l = elements [" \\ ", "λ"]
n' = next n
bs = chooseSize (0,6) (listOf1 genName)
<&> mconcat
next = (`div` 2)
genTy :: Gen Text
genTy = gen 4 where
gen 0 = genCon
gen n = oneof
[ gen 0
-- function types
, conseq [gen n', ws, pure "->", ws, gen n']
-- TODO: type applications (remember precedence lol)
]
where n' = n `div` 2
instance Arbitrary Rec where
arbitrary = elements [Rec,NonRec]
chooseSize :: (Int, Int) -> Gen a -> Gen a
chooseSize (a,b) g = do
n <- chooseInt (a,b)
resize n g
-- | @pprec q p s@ wraps @s@ with parens when @p <= q@
pprec :: (IsString a, Monoid a) => Int -> Int -> a -> a
pprec maxp p
| p <= maxp = id
| otherwise = wrapParens
wrapParens :: (IsString a, Monoid a) => a -> a
wrapParens t = "(" <> t <> ")"
conseq :: (Applicative f, Monoid m, Traversable t)
=> t (f m)
-> f m
conseq tfm = sequenceA tfm <&> the_cool_kid's_concat
-- me when `concat` is generalised in the container but specialised in the
-- value, and `mconcat` is specialised in the container but generalised in
-- the value. shoutout `foldMap id`
where the_cool_kid's_concat = foldMap id
genName :: Gen Name
genName = T.pack <$> liftA2 (:) small namechars where
small = elements ['a'..'z']
genCon :: Gen Name
genCon = T.pack <$> liftA2 (:) large namechars where
large = elements ['A'..'Z']
genSymName :: Gen Name
genSymName = T.pack <$> liftA2 (:) symbol symchars where
symbol = elements nameSymbols
genSymCon :: Gen Name
genSymCon = T.pack . (':' :) <$> symchars
namechars :: Gen String
namechars = liftArbitrary namechar where
namechar :: Gen Char
namechar = elements $ ['a'..'z'] <> ['A'..'Z'] <> ['0'..'9'] <> "'"
nameSymbols :: [Char]
nameSymbols = "!#$%&*+./<=>?@^|-~"
symchars :: Gen String
symchars = liftArbitrary symchar where
symchar = elements $ ':' : nameSymbols
txt :: (IsText t) => t -> Doc
txt t = t ^. unpacked & text
ws :: (IsString a) => Gen a
ws = elements [""," ", " "]
ws1 :: (IsString a) => Gen a
ws1 = elements [" ", " "]
----------------------------------------------------------------------------------
-- | Two bodies of source code are considered congruent iff the parser produces
-- identical ASTs for both.
congruentSrc :: ProgramSrc -> ProgramSrc -> Bool
congruentSrc = (==) `on` (justParseSrc . T.unpack . coerce)
----------------------------------------------------------------------------------
-- TODO: unparseCoreProg :: Program -> [CoreToken]
-- womp womp.
-- TODO: implement shrink
-- | @unparseCoreProg@ should be inverse to @parseCoreProg@ up to source code
-- congruency, newtype coercion and errors handling.
unparseCoreProg :: Program' -> ProgramSrc
unparseCoreProg p = unparseTypeSigs (p ^. programTypeSigs)
<> unparseScDefs (p ^. programScDefs)
unparseTypeSigs :: H.HashMap Name Type -> ProgramSrc
unparseTypeSigs = H.foldrWithKey f mempty
where f k v a = unparseTypeSig k v <> ";\n\n" <> a
unparseTypeSig :: Name -> Type -> ProgramSrc
unparseTypeSig n t = unparseName n <> " :: " <> unparseType t
unparseName :: Name -> ProgramSrc
unparseName n
| T.head n `elem` (':' : nameSymbols) = coerce $ wrapParens n
| otherwise = coerce n
unparseType :: Type -> ProgramSrc
unparseType = go 0 where
go :: Int -> Type -> ProgramSrc
-- (:->) is a special case of TyApp, but we want the infix syntax
go p (a :-> b) = a : assocFun b
<&> go 1
& coerce (T.intercalate " -> ")
& pprec 0 p
go p a@(TyApp f x) = assocApp a
<&> go 1
& coerce (T.intercalate " ")
& pprec 1 p
go _ TyFun = "(->)"
go _ (TyCon a) = unparseName a
go _ (TyVar a) = unparseName a
assocFun :: Type -> [Type]
assocFun (a :-> b) = a : assocFun b
assocFun x = [x]
assocApp :: Type -> [Type]
assocApp (TyApp f x) = assocApp f ++ [x]
assocApp x = [x]
unparseScDefs :: [ScDef'] -> ProgramSrc
unparseScDefs = foldr f mempty where
f sc a = unparseScDef sc <> ";\n\n" <> a
unparseScDef :: ScDef' -> ProgramSrc
unparseScDef (ScDef n as e) = (unparseName <$> (n:as)) <> ["=", unparseExpr e]
& coerce (T.intercalate " ")
unparseExpr :: Expr' -> ProgramSrc
unparseExpr = go 0 where
go :: Int -> Expr' -> ProgramSrc
go _ (Var n) = unparseName n
go _ (Con t a) = mconcat ["Pack{",srcShow t," ",srcShow a,"}"]
go _ (Lit l) = unparseLit l
go p a@(App _ _) = srci " " (go 1 <$> assocApp a)
& pprec 0 p
go p (Lam bs e) = "λ" <> srci " " (unparseName <$> bs)
<> " -> " <> go 0 e
& pprec 0 p
go p (Let r bs e) = mconcat [lw," { ",bs'," } in ",go 0 e]
& pprec 0 p
where
lw = case r of { NonRec -> "let"; Rec -> "letrec" }
bs' = srci "; " $ unparseBinding <$> bs
go p (Case e as) = mconcat ["case ",go 0 e," of {",as',"}"]
& pprec 0 p
where as' = srci "; " (unparseAlter <$> as)
assocApp (App f x) = assocApp f ++ [x]
assocApp f = [f]
srci :: ProgramSrc -> [ProgramSrc] -> ProgramSrc
srci = coerce T.intercalate
unparseBinding :: Binding' -> ProgramSrc
unparseBinding (k := v) = mconcat [unparseName k, " = ", unparseExpr v]
unparseLit :: Lit -> ProgramSrc
unparseLit (IntL n) = srcShow n
srcShow :: (Show a) => a -> ProgramSrc
srcShow = coerce . T.pack . show
unparseAlter :: Alter' -> ProgramSrc
unparseAlter (Alter (AltData t) as e) = srcShow t <> " " <> coerce (T.unwords as)
<> " -> " <> unparseExpr e