~jack/misc

ref: 792d1a6fb90d4bd33438b6e5e72a62737f8589d3 misc/stg/src/Language/STG/EvalApplyAST.hs -rw-r--r-- 4.4 KiB
792d1a6fJack Kelly stg: implement the eval/apply machine a month ago
                                                                                
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedLists #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}

-- | AST for Eval/Apply variant of STG, in the style of
-- https://simonmar.github.io/bib/papers/eval-apply.pdf
module Language.STG.EvalApplyAST where

import Control.Lens.TH
import Relude hiding (Alt)

type Prog = Binds

type Binds = Map Var Object

type Con = Text

type Var = Text

data Literal = I Int | D Double deriving (Eq, Show)

data Atom = L Literal | V Var deriving (Eq, Show)

data Arity = Unknown | Known Int deriving (Eq, Show)

data Alts
  = ConAlts [ConAlt] (Maybe DefAlt)
  | PrimAlt DefAlt
  deriving (Eq, Show)

data ConAlt = ConPat Con [Var] Expr deriving (Eq, Show)

data DefAlt = Default Var Expr deriving (Eq, Show)

data PrimOp = IAdd | ISub | IMul | DAdd | DSub | DMul | DDiv
  deriving (Eq, Show)

data Object
  = Fun [Var] Expr
  | Pap Var [Atom]
  | Con Con [Atom]
  | Thunk Expr
  | Blackhole
  deriving (Eq, Show)

data Expr
  = EAtom Atom
  | EFunCall Var Arity [Atom]
  | EPrimOp PrimOp [Atom]
  | ELet Var Object Expr
  | ECase Expr Alts
  deriving (Eq, Show)

$(makePrisms ''Atom)
$(makePrisms ''Object)
$(makePrisms ''Expr)

stgSub :: Var -> Atom -> Expr -> Expr
stgSub v a e =
  let subArgs :: [Atom] -> [Atom]
      subArgs = map $ \arg -> case arg of
        V v' | v == v' -> a
        _ -> arg
   in case e of
        EAtom (V v') | v == v' -> EAtom a
        EAtom _ -> e
        EFunCall f ar args -> EFunCall f' ar $ subArgs args
          where
            f' =
              if v == f
                then case a of
                  V fv -> fv
                  L _ -> error "cannot sub literal over a function call"
                else f
        EPrimOp p args -> EPrimOp p $ subArgs args
        ELet x o letBody
          | x == v -> e
          | otherwise -> ELet x o' $ stgSub v a letBody
          where
            o' = case o of
              Fun args body
                | v `elem` args -> o
                | otherwise -> Fun args $ stgSub v a body
              Pap {} -> o
              Con c args -> Con c $ subArgs args
              Thunk expr -> Thunk $ stgSub v a expr
              Blackhole -> Blackhole
        ECase c alts -> ECase (stgSub v a c) alts'
          where
            alts' = case alts of
              ConAlts cas def ->
                let mapConAlts = map $ \alt -> case alt of
                      ConPat con args body
                        | v `elem` args -> alt
                        | otherwise -> ConPat con args $ stgSub v a body
                 in ConAlts (mapConAlts cas) $ mapDefAlt <$> def
              PrimAlt def -> PrimAlt $ mapDefAlt def

            mapDefAlt d@(Default v' body)
              | v == v' = d
              | otherwise = Default v' $ stgSub v a body

objectIsValue :: Object -> Bool
objectIsValue = \case
  Fun {} -> True
  Pap {} -> True
  Con {} -> True
  Thunk {} -> False
  Blackhole -> False

stgProg :: Prog
stgProg =
  [ ("intMul", stgIntMul),
    ("main", stgMain),
    ("map", stgMap),
    ("nil", stgNil),
    ("one", Con "I" [L (I 1)]),
    ("two", Con "I" [L (I 2)]),
    ("three", Con "I" [L (I 3)])
  ]

stgIntMul :: Object
stgIntMul =
  Fun ["x", "y"]
    . ECase (EAtom $ V "x")
    $ ConAlts
      [ ConPat "I" ["x#"]
          . ECase (EAtom $ V "y")
          $ ConAlts
            [ ConPat "I" ["y#"]
                . ECase (EPrimOp IMul [V "x#", V "y#"])
                . PrimAlt
                . Default "r#"
                . ELet "r" (Con "I" [V "r#"])
                . EAtom
                $ V "r"
            ]
            Nothing
      ]
      Nothing

stgMain :: Object
stgMain =
  Thunk $
    ELet
      "double"
      (Thunk $ EFunCall "intMul" (Known 2) [V "two"])
      . ELet "t3" (Con "Nil" [])
      . ELet "t2" (Con "Cons" [V "three", V "t3"])
      . ELet "t1" (Con "Cons" [V "two", V "t2"])
      . ELet "list" (Con "Cons" [V "one", V "t1"])
      $ EFunCall "map" (Known 2) [V "double", V "list"]

stgMap :: Object
stgMap =
  Fun ["f", "xs"] $
    ECase
      (EAtom (V "xs"))
      ( ConAlts
          [ ConPat "Nil" [] . EAtom $ V "nil",
            ConPat "Cons" ["y", "ys"]
              . ELet "h" (Thunk $ EFunCall "f" Unknown [V "y"])
              . ELet "t" (Thunk $ EFunCall "map" (Known 2) [V "f", V "ys"])
              . ELet "r" (Con "Cons" [V "h", V "t"])
              . EAtom
              $ V "r"
          ]
          Nothing
      )

stgNil :: Object
stgNil = Con "Nil" []