~jack/misc

ref: 792d1a6fb90d4bd33438b6e5e72a62737f8589d3 misc/stg/src/Language/STG/EAMachine.hs -rw-r--r-- 5.8 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
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TemplateHaskell #-}

module Language.STG.EAMachine where

import Control.Lens hiding (uncons)
import Language.STG.EvalApplyAST
import Relude hiding (Alt)
import qualified Text.Show as S

data S = S
  { expr :: Expr,
    stack :: [Cont],
    heap :: Heap,
    gensyms :: Gensyms
  }
  deriving (Eq, Show)

type Heap = Map Var Object

data Cont
  = CCase Alts
  | CUpdate Var
  | CApply [Atom]
  deriving (Eq, Show)

newtype Gensyms = Gensyms [Var]

instance Show Gensyms where
  show _ = "Gensyms"

instance Eq Gensyms where
  Gensyms xs == Gensyms ys = viaNonEmpty head xs == viaNonEmpty head ys

$(makePrisms ''Cont)

gNext :: Gensyms -> (Gensyms, Var)
gNext (Gensyms (v : vs)) = (Gensyms vs, v)
gNext _ = error "gNext: impossible"

start :: Prog -> S
start p =
  S
    { expr = EAtom $ V "main",
      stack = [],
      heap = p,
      gensyms = Gensyms $ ("_g" <>) . show <$> ([0 ..] :: [Int])
    }

step :: S -> S
step s = case catMaybes $ sequenceA rules s of
  [s'] -> s'
  ss ->
    error $
      "Should only have one matching rule (got: " <> show (length ss) <> ")"
  where
    rules =
      [ rLet,
        rCaseConAny,
        rCase,
        rRet,
        rThunk,
        rUpdate,
        rKnownCall,
        rPrimOp,
        rExact,
        rCallKPap2,
        rTcall,
        rPcall,
        rRetFun
      ]

rCase :: S -> Maybe S
rCase s@S {..} = do
  (e, alts) <- expr ^? _ECase
  guard $ case e of
    EAtom a -> not $ case a of
      L {} -> True
      V v -> heap ^? ix v . to objectIsValue == Just True
    _ -> True
  Just $
    s
      { expr = e,
        stack = CCase alts : stack
      }

rCaseConAny :: S -> Maybe S
rCaseConAny s@S {..} = do
  (EAtom a, alts) <- expr ^? _ECase
  guard $ case a of
    L {} -> True
    V v -> heap ^? ix v . to objectIsValue == Just True
  case (a, alts) of
    (L {}, PrimAlt (Default x e)) ->
      Just $ s {expr = stgSub x a e}
    (V var, ConAlts cas def) ->
      do
        val <- heap ^? ix var
        case val of
          Con c args ->
            let loop :: [ConAlt] -> Maybe S
                loop [] =
                  def <&> \(Default v body) ->
                    s {expr = stgSub v a body}
                loop (ConPat con vars body : cs)
                  | c == con =
                    Just $
                      s
                        { expr = foldl' (flip $ uncurry stgSub) body $ zip vars args
                        }
                  | otherwise = loop cs
             in loop cas
          _ -> def <&> \(Default v body) -> s {expr = stgSub v a body}
    _ -> Nothing

rRet :: S -> Maybe S
rRet s@S {..} = do
  a <- expr ^? _EAtom
  guard $ case a of
    L {} -> True
    V v -> heap ^? ix v . to objectIsValue == Just True
  (CCase alts, s') <- uncons stack
  Just $
    s
      { expr = ECase (EAtom a) alts,
        stack = s'
      }

rLet :: S -> Maybe S
rLet s@S {..} = do
  (x, o, e) <- expr ^? _ELet
  let (g', x') = gNext gensyms
  Just $
    s
      { expr = stgSub x (V x') e,
        heap = heap & at x' ?~ o,
        gensyms = g'
      }

rThunk :: S -> Maybe S
rThunk s@S {..} = do
  v <- expr ^? _EAtom . _V
  e <- heap ^? ix v . _Thunk
  Just $
    s
      { expr = e,
        stack = CUpdate v : stack,
        heap = heap & at v ?~ Blackhole
      }

rUpdate :: S -> Maybe S
rUpdate s@S {..} = do
  V y <- expr ^? _EAtom
  (CUpdate x, s') <- uncons stack
  hy <- heap ^? ix y
  guard $ objectIsValue hy
  Just $ s {stack = s', heap = heap & at x ?~ hy}

rKnownCall :: S -> Maybe S
rKnownCall s@S {..} = do
  (f, Known n, args) <- expr ^? _EFunCall
  guard $ n == length args
  (vars, body) <- heap ^? ix f . _Fun
  Just $ s {expr = foldl' (flip $ uncurry stgSub) body $ zip vars args}

rPrimOp :: S -> Maybe S
rPrimOp s@S {..} = do
  (p, args) <- expr ^? _EPrimOp
  case (p, args) of
    (IAdd, [L (I x), L (I y)]) -> Just $ s {expr = EAtom . L . I $ x + y}
    (ISub, [L (I x), L (I y)]) -> Just $ s {expr = EAtom . L . I $ x - y}
    (IMul, [L (I x), L (I y)]) -> Just $ s {expr = EAtom . L . I $ x * y}
    (DAdd, [L (D x), L (D y)]) -> Just $ s {expr = EAtom . L . D $ x + y}
    (DSub, [L (D x), L (D y)]) -> Just $ s {expr = EAtom . L . D $ x - y}
    (DMul, [L (D x), L (D y)]) -> Just $ s {expr = EAtom . L . D $ x * y}
    (DDiv, [L (D x), L (D y)]) -> Just $ s {expr = EAtom . L . D $ x / y}
    _ -> Nothing

rExact :: S -> Maybe S
rExact s@S {..} = do
  (f, Unknown, args) <- expr ^? _EFunCall
  (vars, body) <- heap ^? ix f . _Fun
  guard $ length vars == length args
  Just $ s {expr = foldl' (flip $ uncurry stgSub) body $ zip vars args}

rCallKPap2 :: S -> Maybe S
rCallKPap2 s@S {..} = do
  (f, Known _, args) <- expr ^? _EFunCall
  (vars, body) <- heap ^? ix f . _Fun
  case compare (length args) (length vars) of
    EQ -> Nothing
    LT ->
      let (g', p) = gNext gensyms
       in Just $
            s
              { expr = EAtom (V p),
                heap = heap & at p ?~ Pap f args,
                gensyms = g'
              }
    GT ->
      let (applied, leftover) = splitAt (length vars) args
       in Just $
            s
              { expr = foldl' (flip $ uncurry stgSub) body $ zip vars applied,
                stack = CApply leftover : stack
              }

rTcall :: S -> Maybe S
rTcall s@S {..} = do
  (f, Unknown, args) <- expr ^? _EFunCall
  e <- heap ^? ix f . _Thunk
  Just $ s {expr = e, stack = CApply args : stack}

rPcall :: S -> Maybe S
rPcall s@S {..} = do
  (f, Unknown, args) <- expr ^? _EFunCall
  (g, rest) <- heap ^? ix f . _Pap
  Just $ s {expr = EFunCall g Unknown $ args ++ rest}

rRetFun :: S -> Maybe S
rRetFun s@S {..} = do
  V f <- expr ^? _EAtom
  hf <- heap ^? ix f
  guard $ case hf of
    Fun {} -> True
    Pap {} -> True
    _ -> False
  (CApply rest, s') <- uncons stack
  Just $ s {expr = EFunCall f Unknown rest, stack = s'}