higher-order defun

2026-07-12 20:58:10 -06:00
parent 4522e455dd
commit 269d956566
8 changed files with 194 additions and 46 deletions
@@ -6,8 +6,7 @@
 {-# LANGUAGE OverloadedLists #-}
 {-# OPTIONS_GHC -Wno-incomplete-patterns #-}
 module Gyehoek.CPS.Lower
-  (
-  lower, lowerProgram) where
+  (lower, lowerProgram) where

 import Gyehoek.CPS.Syntax
 import Data.Generics.Labels
@@ -31,15 +30,18 @@ import qualified Data.Vector.Strict as V
 import Data.IntMap.Strict (IntMap)
 import Data.String.Interpolate
 import Gyehoek.Wasm qualified as Wasm
-import Gyehoek.Wasm (i32, ins, sxp, eq, ref, i31, Type (..))
+import Gyehoek.Wasm (i32, ins, sxp, eq, ref, i31, Type (..), Idx, GenMod)
 import Language.Sexp.Located (pattern ParenList)


-data Env = MkEnv { vars :: Vector Name }
+data Env = MkEnv
+  { runtime :: Runtime
+  , vars :: Vector Name
+  }
  deriving (Show, Generic)

 emptyEnv :: Env
-emptyEnv = MkEnv mempty
+emptyEnv = MkEnv (error "fuck") mempty

 type instance Index Env = Natural
 type instance IxValue Env = Name
@@ -47,15 +49,31 @@ type instance IxValue Env = Name
 instance Ixed Env where
  ix i = #vars . ix (fromIntegral i)

+data Runtime = MkRuntime
+  { consIdx :: Idx
+  }
+  deriving (Show, Generic)
+
+
+
+-- | @makeSmallFixnum@ emits an expression injecting the i32 on top
+-- of the stack into the SCM unitype.
+makeSmallFixnum :: Wasm.Expr
+makeSmallFixnum = mconcat
+  [ ins "i32.const" [sxp @Int 2]
+  , ins "i32.shl" []
+  , ins "ref.i31" []
+  ]
+


 lowerVal :: Env -> Val -> Wasm.Expr

 lowerVal g (ValLit l) =
  case l of
-    LitInt n ->
+    LitInt n -> 
      ins "i32.const" [sxp n]
-      <> ins "ref.i31" []
+      <> makeSmallFixnum
    LitBool b ->
      ins "i32.const" [sxp @Int $ if b then 1 else 0]
      <> ins "ref.i31" []
@@ -65,9 +83,9 @@ lowerVal g (ValVar x) = ins "local.get" [sxp l]
  where
    l = V.elemIndex x g.vars ^?! _Just

-lower' :: Env -> Exp -> Wasm.Expr
+lower' :: (GenMod :> es) => Env -> Exp -> Eff es Wasm.Expr

-lower' g (Halt [e]) = lowerVal g e
+lower' g (Halt [e]) = pure $ lowerVal g e

 lower' g (ExpPrim p rs e) =
  case p of
@@ -76,26 +94,29 @@ lower' g (ExpPrim p rs e) =
  where
    r = head rs

-lower' g (ExpIf c t f) =
-  lowerVal g c
-  <> Wasm.if' (Wasm.result [i32])
-    (lower' g t)
-    (lower' g f)
+lower' g (ExpIf c t f) = do
+  t' <- lower' g t
+  f' <- lower' g f
+  pure $ lowerVal g c
+    <> Wasm.if' (Wasm.result [i32]) t' f'

 lowerBinOp
-  :: _
-  -> _ -> _ -> _ -> _ -> _ -> Wasm.Expr
-lowerBinOp op g x y r e =
-  lowerVal g x
-  <> ins "ref.cast" [sxp $ ref i31]
-  <> ins "i31.get_s" []
-  <> lowerVal g y
-  <> ins "ref.cast" [sxp $ ref i31]
-  <> ins "i31.get_s" []
-  <> ins op []
-  <> ins "ref.i31" []
-  <> ins "local.set" [sxp n]
-  <> lower' g' e
+  :: (GenMod :> es)
+  => Text -> Env -> Val -> Val -> Name -> Exp -> Eff es Wasm.Expr
+lowerBinOp op g x y r e = do
+  e' <- lower' g' e
+  pure . mconcat $
+    [ lowerVal g x
+    , ins "ref.cast" [sxp $ ref i31]
+    , ins "i31.get_s" []
+    , lowerVal g y
+    , ins "ref.cast" [sxp $ ref i31]
+    , ins "i31.get_s" []
+    , ins op []
+    , ins "ref.i31" []
+    , ins "local.set" [sxp n]
+    , e'
+    ]
  where
    g' = g & #vars <>~ [r]
    n = length (g ^. #vars)
@@ -106,10 +127,17 @@ scm = ref eq



-lower :: Exp -> Eff es Text
-lower e = fmap Wasm.renderModule . Wasm.execGenMod $ do
+emitRuntime :: GenMod :> es => Eff es Runtime
+emitRuntime = do
  Wasm.deftypeNamed "$heap-object" $ Wasm.sub [] $ Wasm.struct
    [ Wasm.mut i32 ]
+  consIdx <- Wasm.defun _ _ _ _
+  pure $ MkRuntime {consIdx}
+
+lower :: Exp -> Eff es Text
+lower e = fmap Wasm.renderModule . Wasm.execGenMod $ do
+  -- runtime <- emitRuntime
+  let env = MkEnv _runtime mempty
  main <- Wasm.defun [] [scm] [scm, scm, scm, scm, scm] \_ ->
    lower' emptyEnv e
  Wasm.export "main" "func" main
@@ -48,7 +48,7 @@ import GHC.IO.Unsafe (unsafePerformIO)
 import qualified Data.Text.IO as TIO
 import Control.Monad (join)
 import qualified Language.Sexp.Located as SexpLoc
-import Data.Void (absurd)
+import Data.Void (absurd, Void)
 import Data.Coerce (coerce)
 import qualified Data.Map

@@ -3,6 +3,7 @@
 {-# LANGUAGE DeepSubsumption #-}
 {-# LANGUAGE NoFieldSelectors #-}
 {-# LANGUAGE OverloadedRecordDot #-}
+{-# LANGUAGE RecordPuns #-}
 {-# LANGUAGE DuplicateRecordFields #-}
 {-# LANGUAGE QuasiQuotes #-}
 {-# LANGUAGE OverloadedLabels #-}
@@ -22,6 +23,7 @@ module Gyehoek.Wasm
  , Expr
  , Instr
  , GenMod
+  , Idx
  , i32
  , export
  , ins
@@ -125,7 +127,8 @@ data Idx

 data GenMod :: Effect where
  DefRecType :: List Type -> GenMod m (List Idx)
-  Defun :: List Type -> List Type -> List Type -> (Idx -> Expr) -> GenMod m Idx
+  Defun :: List Type -> List Type -> List Type
+        -> (Idx -> m Expr) -> GenMod m Idx
  Start :: Idx -> GenMod m ()
  Export :: Text -> Text -> Idx -> GenMod m ()

@@ -151,9 +154,9 @@ deftypeNamed name (MkType t) = void $ send (DefRecType [namedType name t])
 defun
  :: (GenMod :> es)
  => List Type -> List Type -> List Type
-  -> (Idx -> Expr)
+  -> (Idx -> Eff es Expr)
  -> Eff es Idx
-defun params result locals code = send $ Defun params result locals code
+defun params res locals code = send $ Defun params res locals code

 -- defun
 --   :: (GenMod :> es)
@@ -163,7 +166,7 @@ defun params result locals code = send $ Defun params result locals code
 -- defun params result locals code =
 --   send $ Defun params result locals (runPureEff . execWriterLocal . code)

-runGenMod :: Eff (GenMod : es) a -> Eff es (a, Module)
+runGenMod :: forall es a. Eff (GenMod : es) a -> Eff es (a, Module)
 runGenMod = 
  reinterpret (runStateLocal (mempty :: Module)) \cases
    _ (DefRecType ts) -> state \m ->
@@ -176,12 +179,18 @@ runGenMod =
       #exports <>= V.singleton e
      where e = MkExport $ ParenList
              [ "export", sxp name, ParenList [ "func", sxp idx ] ]
-    _ (Defun params res locals code) -> state \m ->
-      let idx = IdxNumeric . fromIntegral . length $ m.functions
-      in ( idx
-         , m & #functions <>~ V.singleton
-             (MkFunction params res locals (code idx))
-         )
+    env (Defun params result locals code) ->
+      localSeqUnlift env \unlift ->
+        stateM \m -> do
+          -- the least unused function index, computed as the number
+          -- of currently allocated functions.
+          let idx = IdxNumeric . fromIntegral . length $ m.functions
+          -- the body is computed with access to the newly allocated
+          -- index `idx` for the sake of recursive occurences.
+          body <- unlift $ code idx
+          let func = MkFunction {params,result,locals,body}
+          let m' = m & #functions <>~ V.singleton func
+          pure (idx, m')

 execGenMod = fmap snd . runGenMod

@@ -214,7 +223,7 @@ i31 = MkType $ Symbol "i31"
 instance SexpIso Idx where
  sexpIso = match
    $ With (\numeric -> num >>> numeric)
-    $ With (\named -> symbol >>> named)
+    $ With (\named -> name >>> named)
    $ End
    where
      num = Sexp.integer >>> Sexp.partialOsi f g
@@ -223,6 +232,11 @@ instance SexpIso Idx where
                            <> Sexp.expected "natural"
              | otherwise = Right $ fromIntegral n
          g n = fromIntegral n
+      l :: Prism' Text Text
+      l = prefixed "$"
+      name = Sexp.symbol >>> Sexp.partialOsi
+        (maybe (Left $ Sexp.expected "$-prefixed sym") Right . preview l)
+        (review l)

 instance SexpIso RecType where
  sexpIso = with \rectype ->
@@ -287,8 +301,8 @@ instance SexpIso Sexp where
 instance Each Expr Expr Instr Instr where
  each = #MkExpr . each

-sxp :: SexpIso a => a -> Sexp
-sxp e = Sexp.toSexp sexpIso e ^?! _Right
+sxp :: HasCallStack => SexpIso a => a -> Sexp
+sxp e = either error id . Sexp.toSexp sexpIso $ e

 ins :: Text -> List Sexp -> Expr
 ins op [] = [ MkInstr $ Symbol op ]