Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
311 commits
Select commit Hold shift + click to select a range
f1a209b
remove Label from Block.mls
ChingLongTin Nov 3, 2025
55f4fa1
rename helpers
ChingLongTin Nov 3, 2025
124232f
update Block.mls
ChingLongTin Nov 3, 2025
7c8f3d9
adapt NeilKleistGao's example code
ChingLongTin Nov 3, 2025
6997e89
rename helpers
ChingLongTin Nov 3, 2025
36abad1
refactor ruleMatch
ChingLongTin Nov 3, 2025
64960a8
fix ruleTup
ChingLongTin Nov 3, 2025
a264406
reduce supported cases to make mvp
ChingLongTin Nov 3, 2025
4ae8897
Update Syntax.mls
ChingLongTin Nov 1, 2025
88c3a5e
move detecting staged annotation to ClsLikeDefn
ChingLongTin Nov 4, 2025
17701b9
move importing to diff testing compiler flag
ChingLongTin Nov 4, 2025
5ccb64c
revert formatting changes
ChingLongTin Nov 4, 2025
980d89c
Update Printer.scala
ChingLongTin Nov 5, 2025
ce1e38c
Update Lowering.scala
ChingLongTin Nov 5, 2025
ec3e7b8
Merge branch 'hkmc2' into staged-syntax
ChingLongTin Nov 5, 2025
317c6f6
remove redundant :pt
ChingLongTin Nov 5, 2025
9023eec
fix syntax error
ChingLongTin Nov 5, 2025
9fc83dc
fix type hint for arms in Match
ChingLongTin Nov 5, 2025
7b702f3
add preliminary printCode function
yeungsinchun Nov 5, 2025
e4d7b22
move instrumentation to JSBackendDiffMaker
ChingLongTin Nov 6, 2025
6576f28
rename compiler flag to :staging
ChingLongTin Nov 6, 2025
4191464
Revert "remove redundant :pt"
ChingLongTin Nov 6, 2025
817e828
combine import statements to list
ChingLongTin Nov 6, 2025
e4f947d
Merge branch 'staged-syntax' into instrument
ChingLongTin Nov 6, 2025
33a3749
reword function
ChingLongTin Nov 6, 2025
b10e975
Revert "move instrumentation to JSBackendDiffMaker"
ChingLongTin Nov 6, 2025
b93b5eb
Merge branch 'staged-syntax' into instrument-simple
ChingLongTin Nov 6, 2025
791b5b4
prune implementation for mvp
ChingLongTin Nov 6, 2025
578dfd1
implement rudimentary instrumentation logic
ChingLongTin Nov 6, 2025
27bf507
implement printing for DynSelect
ChingLongTin Nov 6, 2025
9c975f6
only stage functions in staged modules
ChingLongTin Nov 6, 2025
ff6b121
update Syntax.mls
ChingLongTin Nov 6, 2025
068fd37
remove bugs in pretty printing
yeungsinchun Nov 7, 2025
f70f346
Merge branch 'pretty-print' into instrument-simple
yeungsinchun Nov 7, 2025
3f32b39
formatting
ChingLongTin Nov 8, 2025
d3bd300
add back Symbol class
ChingLongTin Nov 8, 2025
8ef91e9
reuse original function name
ChingLongTin Nov 8, 2025
3c7f58d
update test cases
ChingLongTin Nov 8, 2025
2ecda34
fix ruleVal
ChingLongTin Nov 8, 2025
9ebf26b
add some debug symbol names
ChingLongTin Nov 8, 2025
0198eb9
rename PathLike to ArgWrappable
ChingLongTin Nov 8, 2025
8fa331b
fix applyBlock impl
ChingLongTin Nov 9, 2025
e167b20
refactor transformFunDefn
ChingLongTin Nov 10, 2025
0e08e07
remove spread from Arg
ChingLongTin Nov 10, 2025
827f6af
clean up outdated comments and helpers
ChingLongTin Nov 10, 2025
57b45fd
implement ruleTup
ChingLongTin Nov 10, 2025
37eb7a1
formatting
ChingLongTin Nov 10, 2025
69f5997
implement ruleSel and ruleInst
ChingLongTin Nov 10, 2025
6831c61
remove outdated comment
ChingLongTin Nov 10, 2025
ab92537
clean up symName printing
ChingLongTin Nov 10, 2025
29c1109
formatting
ChingLongTin Nov 10, 2025
55fdf26
Add shape in diff testing
yeungsinchun Nov 11, 2025
1e0eb65
update note
ChingLongTin Nov 17, 2025
b93a91c
update PrintCode printout
ChingLongTin Nov 17, 2025
772c93a
update Shape constructors
ChingLongTin Nov 17, 2025
fbe45ea
add context
ChingLongTin Nov 17, 2025
873c8cc
remove outdated comments
ChingLongTin Nov 17, 2025
31809c7
correc comment
yeungsinchun Nov 20, 2025
91e5197
fix pretty print to print mlscript
yeungsinchun Nov 20, 2025
f1886ec
refactor
ChingLongTin Nov 20, 2025
d99eda9
follow JSBuilder pretty printing for DynSelect in Printer
ChingLongTin Nov 20, 2025
02303a3
remove Field from Case
ChingLongTin Nov 20, 2025
d6dc21e
change transformBlock signature
ChingLongTin Nov 20, 2025
cde7a68
implement more rules
ChingLongTin Nov 20, 2025
1f3ee7c
remove import Option from Shape.mls
ChingLongTin Nov 20, 2025
0499166
replace Multiple with ShapeSet
ChingLongTin Nov 20, 2025
fbb869c
misc fixes to Shape.mls
ChingLongTin Nov 20, 2025
adf5590
patch showLiteral for null and undefined
ChingLongTin Nov 20, 2025
ada9c3a
replace ruleVal
ChingLongTin Nov 21, 2025
9f6b1ab
clean up
ChingLongTin Nov 21, 2025
ba80ea4
add symbol to Class Shape for matching
ChingLongTin Nov 21, 2025
8efd61d
add spread field to Arg
ChingLongTin Nov 21, 2025
e19dbd6
update PrintCode.mls
ChingLongTin Nov 21, 2025
146a198
add ruleCls
ChingLongTin Nov 21, 2025
e1aedc7
formatting
ChingLongTin Nov 21, 2025
ed85e75
progress on rest of the rules
ChingLongTin Nov 21, 2025
53a0aa4
replace Block.Case with Pattern
ChingLongTin Nov 21, 2025
3a992ad
implement ShapeSet and ShapeMap
ChingLongTin Nov 21, 2025
5840db4
fix equality checks for Symbol
ChingLongTin Nov 22, 2025
bc2f624
implement Shape{Map, Set}
ChingLongTin Nov 22, 2025
e2454b4
fill out ShapeSet functionailty
ChingLongTin Nov 22, 2025
746cbd3
link ShapeSet symbols for instrumentation
ChingLongTin Nov 22, 2025
9915656
fix Shape, ShapeSet implementations
ChingLongTin Nov 22, 2025
4930252
remove TrivialResult
ChingLongTin Nov 22, 2025
b3c707f
replace Shape with ShapeSet during instrumentation
ChingLongTin Nov 23, 2025
e76c1fd
add message to assertions
ChingLongTin Nov 23, 2025
d792553
refactor match rules
ChingLongTin Nov 23, 2025
1020d38
add debug print Shape
ChingLongTin Nov 23, 2025
eece7ac
print more information for ClsLikeDefn
ChingLongTin Nov 23, 2025
df68b1d
combine Pattern with Case
ChingLongTin Nov 23, 2025
aa3b19a
fill out more match instrumentation
ChingLongTin Nov 23, 2025
4af5839
update comments
ChingLongTin Nov 23, 2025
2c67c8e
update PrintCode.mls
ChingLongTin Nov 23, 2025
c90118c
move optionNme inside stageCode scope
ChingLongTin Nov 23, 2025
f149026
Merge remote-tracking branch 'mlscript/hkmc2' into instrument
ChingLongTin Nov 23, 2025
aae0086
Revert "move optionNme inside stageCode scope"
ChingLongTin Nov 23, 2025
986ace7
use more meaningful names in Functions.mls
ChingLongTin Nov 23, 2025
f0456e1
refactor collectApply
ChingLongTin Nov 23, 2025
d510452
formatting
ChingLongTin Nov 23, 2025
0b2b7ea
update comments and debug info
ChingLongTin Nov 23, 2025
b6c86e1
remove redundant import
ChingLongTin Nov 24, 2025
869c519
match pretty printing syntax for DynSelect with formalization
ChingLongTin Nov 24, 2025
cff8896
fix mkClass
ChingLongTin Nov 24, 2025
0d95197
Merge branch 'hkmc2' into instrument
ChingLongTin Nov 24, 2025
8853371
formatting
ChingLongTin Nov 24, 2025
60c777c
Merge branch 'instrument' of https://github.com/ChingLongTin/mlscript…
ChingLongTin Nov 24, 2025
a343c36
fix syntax errors
ChingLongTin Nov 24, 2025
3c4faec
patch in instrumentation
ChingLongTin Nov 24, 2025
b7c3698
unreserve Option in Elaborator
ChingLongTin Nov 24, 2025
c149a11
remove argument names from class shape
ChingLongTin Nov 24, 2025
1f5c8d7
move functions
ChingLongTin Nov 24, 2025
a6b0162
add ClassSymbol
ChingLongTin Nov 24, 2025
c9f3c6a
formatting
ChingLongTin Nov 24, 2025
db7395d
clean up
ChingLongTin Nov 24, 2025
af3f436
Revert "move functions"
ChingLongTin Nov 24, 2025
89f414c
update diff output
yeungsinchun Nov 25, 2025
6db62ea
pass staged annotation in lowering
yeungsinchun Nov 26, 2025
c1f00e8
Revert "patch in instrumentation"
ChingLongTin Nov 26, 2025
143ce9b
rename functions
ChingLongTin Nov 24, 2025
fcbe8c9
add debug information to rules
ChingLongTin Nov 24, 2025
87b5457
move paramsOpt into ClassSymbol
ChingLongTin Nov 25, 2025
67a099c
implement pruning of bad arms from Match statement
ChingLongTin Nov 26, 2025
9f148af
add debug name to transformSymbol
ChingLongTin Nov 26, 2025
cd26ae1
remove code duplication from ruleBranch
ChingLongTin Nov 26, 2025
2b5b30a
fix typo
ChingLongTin Nov 26, 2025
8b459b3
Merge branch 'fix-pretty-print' into instrument
yeungsinchun Nov 26, 2025
5d63091
update pretty printing to print actual mlscript
yeungsinchun Nov 26, 2025
cc24cae
only pass the "staged" annotation for non-companion module
yeungsinchun Nov 26, 2025
2b1d157
Merge remote-tracking branch 'origin/instrument' into instrument
yeungsinchun Nov 26, 2025
bcafb98
reduce test cases
yeungsinchun Nov 26, 2025
b9eee38
fix class definition printing
yeungsinchun Nov 26, 2025
50f07ca
add match error on pretty printing
yeungsinchun Nov 26, 2025
678dbc9
fix test cases
yeungsinchun Nov 26, 2025
71a8918
some unimportant change
yeungsinchun Nov 26, 2025
7151735
some more unimportant change
yeungsinchun Nov 26, 2025
2548c78
Merge branch 'instrument-with-pretty-print' into instrument
yeungsinchun Nov 26, 2025
f582fa5
printed correctly now
yeungsinchun Nov 26, 2025
0a6087f
Revert "printed correctly now"
yeungsinchun Nov 26, 2025
cd7f7bc
Revert "Merge branch 'instrument-with-pretty-print' into instrument"
yeungsinchun Nov 26, 2025
34d0914
Revert "some more unimportant change"
yeungsinchun Nov 26, 2025
a3925e5
Revert "Merge remote-tracking branch 'origin/instrument' into instrum…
yeungsinchun Nov 26, 2025
79c90ee
Revert "update pretty printing to print actual mlscript"
yeungsinchun Nov 26, 2025
430d941
Merge branch 'hkmc2' into instrument
ChingLongTin Nov 26, 2025
ff58a1a
Reapply "Merge remote-tracking branch 'origin/instrument' into instru…
ChingLongTin Nov 27, 2025
50b6fd5
fix static impl
ChingLongTin Nov 27, 2025
087c960
add inf back to Tup
ChingLongTin Nov 26, 2025
302e310
fix showArg printing
ChingLongTin Nov 26, 2025
4d8e46c
add more Tup test cases
ChingLongTin Nov 26, 2025
f3ee1c6
formatting
ChingLongTin Nov 26, 2025
c703795
fix params type in FunDefn
ChingLongTin Nov 27, 2025
c124135
add Context to ruleEnd
ChingLongTin Nov 28, 2025
45b7b5e
uppdate Match rules
ChingLongTin Nov 28, 2025
408733a
mark pretty printing equality as fixme
ChingLongTin Nov 28, 2025
06312fb
Merge branch 'hkmc2' into instrument
ChingLongTin Dec 11, 2025
0e2fec8
progress: deprecate ShapeSet from StagedPath
ChingLongTin Dec 12, 2025
97ede40
progress: strip shape logic from rules
ChingLongTin Dec 12, 2025
d507d3f
remove inferring shape from Instrumentation
ChingLongTin Dec 12, 2025
ff7385d
progress: add StagedPath to ArgWrappable
ChingLongTin Dec 12, 2025
6db1ee7
progress: simplify impl
ChingLongTin Dec 12, 2025
977433a
progress: formatting
ChingLongTin Dec 12, 2025
b79613a
Merge branch 'hkmc2' into instrument
ChingLongTin Dec 12, 2025
f58ecc5
Merge remote-tracking branch 'mlscript/hkmc2' into instrument
ChingLongTin Dec 18, 2025
e21bff2
fix printCode syntax
ChingLongTin Dec 18, 2025
55e1075
remove equality hack in Shape.mls
ChingLongTin Dec 18, 2025
aa72ded
switch from ValDefn to Scoped
ChingLongTin Dec 18, 2025
326ac69
fix: make new dSym when cloning FunDefn
ChingLongTin Dec 18, 2025
77acb04
remove context from ruleEnd
ChingLongTin Dec 18, 2025
ccc950d
add staging for functions with arguments
ChingLongTin Dec 18, 2025
dee92b1
progress: simplify StagedPath
ChingLongTin Dec 18, 2025
e6fd90a
progress: Make blockCtor produce StagedPath
ChingLongTin Dec 18, 2025
59a4af2
use ClassSymbol for ruleApp
ChingLongTin Dec 18, 2025
5adcfa3
formatting
ChingLongTin Dec 18, 2025
a4c8df0
update tests
ChingLongTin Dec 18, 2025
7abddf9
remove StagedPath
ChingLongTin Dec 22, 2025
f0c6f52
formatting
ChingLongTin Dec 22, 2025
a407e01
inline functions
ChingLongTin Dec 22, 2025
e40a1e6
formatting
ChingLongTin Dec 22, 2025
69a142f
formatting
ChingLongTin Dec 22, 2025
6b012b2
remove Shape{Map, Set} definitions
ChingLongTin Dec 22, 2025
8a713b5
Merge branch 'hkmc2' into instrument
ChingLongTin Dec 22, 2025
51e5513
formatting
ChingLongTin Dec 23, 2025
0cc778e
imrpove symbol information for select
ChingLongTin Dec 23, 2025
e3f052b
simplify applyRuleBranch
ChingLongTin Dec 23, 2025
ed36c03
import option file during MLsDiffMaker init
ChingLongTin Dec 23, 2025
b18776f
remove import statement for option in MLsDiffMaker
ChingLongTin Dec 23, 2025
3df66e2
fixup! remove import statement for option in MLsDiffMaker
ChingLongTin Dec 24, 2025
9212c53
include staging tests in diff tests
ChingLongTin Dec 24, 2025
9f2ca50
remove redundant codegen
ChingLongTin Dec 24, 2025
504d21a
Merge branch 'hkmc2' into instrument
ChingLongTin Dec 24, 2025
247a4c9
Merge branch 'hkmc2' into instrument
ChingLongTin Jan 8, 2026
7da3eea
add Block symbol following pr #375
ChingLongTin Jan 8, 2026
7e115d7
formatting
ChingLongTin Jan 8, 2026
9bcdf99
add Scoped, Loop, Break nodes to Block IR
ChingLongTin Jan 8, 2026
8585c5e
Merge branch 'hkmc2' into instrument
ChingLongTin Jan 13, 2026
6f456bb
Merge branch 'hkmc2' into instrument
ChingLongTin Jan 20, 2026
b16fdf4
replace parameters of functions with its symbol instead of undefined
ChingLongTin Jan 14, 2026
b516237
retain scoped symbols after instrumentation
ChingLongTin Jan 21, 2026
0b968ed
update test
ChingLongTin Jan 21, 2026
78e754f
replace calls to staged functions
ChingLongTin Jan 23, 2026
fe08c3b
Wrap function body
ChingLongTin Jan 23, 2026
efb6ae8
move printing into constructor
ChingLongTin Jan 23, 2026
245da6d
remove _gen suffix from instrumented block
ChingLongTin Jan 27, 2026
1c86af0
remove parameters from instrumentation function
ChingLongTin Jan 27, 2026
7a00ee0
Merge branch 'hkmc2' into instrument
ChingLongTin Feb 2, 2026
a977dd8
Merge branch 'hkmc2' into instrument
ChingLongTin Feb 5, 2026
f5459ac
minor nit
ChingLongTin Feb 6, 2026
dbe58e7
rename TempSymbol when staging
ChingLongTin Feb 10, 2026
7167712
fix concat logic
ChingLongTin Feb 10, 2026
cd16d6f
warn shape propogation doesn't cover multiple parameter lists
ChingLongTin Feb 10, 2026
9d9d449
move ctor staging to a function
ChingLongTin Feb 11, 2026
0a61576
WIP: add class staging test
ChingLongTin Feb 27, 2026
546491a
resolve fixmes
ChingLongTin Feb 27, 2026
0fb3e2e
stage class functions
ChingLongTin Feb 27, 2026
00a534e
Merge remote-tracking branch 'mlscript/hkmc2' into instrument
ChingLongTin Feb 28, 2026
9205ad1
Merge remote-tracking branch 'mlscript/hkmc2' into instrument
ChingLongTin Feb 28, 2026
f2a1ad1
update tests
ChingLongTin Feb 28, 2026
5ed178c
Merge branch 'hkmc2' into instrument
ChingLongTin Mar 3, 2026
a49188b
Merge branch 'hkmc2' into instrument
ChingLongTin Mar 3, 2026
b623e7a
Merge branch 'hkmc2' into instrument
LPTK Mar 4, 2026
55a6ba2
Merge remote-tracking branch 'mlscript/hkmc2' into instrument
ChingLongTin Mar 4, 2026
f087dab
Merge remote-tracking branch 'origin/instrument' into instrument
ChingLongTin Mar 4, 2026
01dd129
always enable staging
ChingLongTin Mar 4, 2026
cd23568
Update hkmc2/shared/src/main/scala/hkmc2/codegen/Instrumentation.scala
ChingLongTin Mar 6, 2026
628218e
replace ??? with ErrorReport
ChingLongTin Mar 6, 2026
401ecc3
Merge branch 'instrument' of https://github.com/ChingLongTin/mlscript…
ChingLongTin Mar 6, 2026
cd94b6e
cleanup
ChingLongTin Mar 6, 2026
5170e25
remove some crashes from .get
ChingLongTin Mar 6, 2026
3cb2ebb
rename transformFunDefn
ChingLongTin Mar 6, 2026
6acee71
remove inf parameter form Shape.Arr and Block.Tup
ChingLongTin Mar 6, 2026
196b544
remove dead code
ChingLongTin Mar 6, 2026
feaab6c
Merge remote-tracking branch 'mlscript/hkmc2' into instrument
ChingLongTin Mar 8, 2026
b3bc782
use :todo for spread test case
ChingLongTin Mar 8, 2026
0b3fe40
Merge branch 'hkmc2' into instrument
LPTK Mar 9, 2026
dd059c9
Merge branch 'hkmc2' into instrument
ChingLongTin Mar 10, 2026
6e1e161
Merge branch 'hkmc2' into instrument
LPTK Mar 10, 2026
99f2fc0
Merge branch 'hkmc2' into instrument
LPTK Mar 10, 2026
befdc88
Update difftests
LPTK Mar 10, 2026
e9b2ae0
fix printing bug
ChingLongTin Mar 10, 2026
28e7b98
extend Block IR Symbol information
ChingLongTin Mar 12, 2026
6beb77c
update tests
ChingLongTin Mar 12, 2026
abba1cd
Merge remote-tracking branch 'mlscript/hkmc2' into patch
ChingLongTin Mar 12, 2026
d3d3305
cleanup
ChingLongTin Mar 12, 2026
af31477
remove redundant Ident type
ChingLongTin Mar 12, 2026
0199444
fixup! update tests
ChingLongTin Mar 12, 2026
79b42cd
cleanup
ChingLongTin Mar 13, 2026
15949d0
fix typo
ChingLongTin Mar 13, 2026
f62f6a8
implement fixes
ChingLongTin Mar 13, 2026
d861b58
remove TypedArray branch
ChingLongTin Mar 13, 2026
9402456
Merge branch 'hkmc2' into patch
ChingLongTin Mar 13, 2026
4403675
Merge branch 'hkmc2' into patch
ChingLongTin Mar 14, 2026
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 55 additions & 27 deletions hkmc2/shared/src/main/scala/hkmc2/codegen/Instrumentation.scala
Original file line number Diff line number Diff line change
Expand Up @@ -5,20 +5,23 @@ import utils.*
import hkmc2.Message.MessageContext

import scala.collection.mutable.HashMap
import scala.util.chaining._
import scala.util.chaining.*

import mlscript.utils.*, shorthands.*

import semantics.*
import semantics.Elaborator.State
import semantics.Elaborator.{State, Ctx, ctx}

import syntax.{Literal, Tree}

// it should be possible to cache some common constructions (End, Option) into the context
// this avoids having to rebuild the same shapes everytime they are needed

// transform Block to Block IR so that it can be instrumented in mlscript
class InstrumentationImpl(using State, Raise):
class Instrumentation(using State, Raise, Ctx) extends BlockTransformer(new SymbolSubst()):
// recover `defn` for when `sym.defn` is `None`, when the definition was generated by other compiler passes
val defnMap = HashMap[Symbol, ClsLikeDefn]()

type ArgWrappable = Path | Symbol
type Context = HashMap[Path, Path]
// TODO: there could be a fresh scope per function body, instead of a single one for the entire program
Expand Down Expand Up @@ -91,22 +94,38 @@ class InstrumentationImpl(using State, Raise):

// transformation helpers

def transformSymbol(sym: Symbol, symName: Str = "sym")(k: Path => Block): Block =
// if sym is ClassSymbol, we may need pOpt to link to the path pointing to the value of the symbol
def transformSymbol(sym: Symbol, pOpt: Option[Path] = N, symName: Str = "sym")(k: Path => Block): Block =
sym match
case clsSym: ClassSymbol =>
clsSym.defn match
case S(defn) =>
val name = scope.allocateOrGetName(sym)
transformParamsOpt(defn.paramsOpt): paramsOpt =>
blockCtor("ClassSymbol", Ls(toValue(name), paramsOpt), symName)(k)
case N =>
raise(ErrorReport(msg"Unable to infer parameters from ClassSymbol in staged module, which are necessary to reconstruct class instances." -> sym.toLoc :: Nil))
End()
case t: TermSymbol if t.defn.exists(_.sym.asCls.isDefined) =>
transformSymbol(t.defn.get.sym.asCls.get, symName)(k)
case t: TermSymbol if t.defn.exists(_.sym.asClsOrMod.isDefined) =>
transformSymbol(t.defn.get.sym.asClsOrMod.get, pOpt, symName)(k)
// retain names to built-in functions or function definitions
case t: TermSymbol if t.defn.exists(_.k == syntax.Fun) =>
blockCtor("Symbol", Ls(toValue(sym.nme)), symName)(k)
case _: BuiltinSymbol =>
// retain names to built-in functions
blockCtor("Symbol", Ls(toValue(sym.nme)), symName)(k)
case clsSym: ClassSymbol if ctx.builtins.virtualClasses(clsSym) =>
blockCtor("VirtualClassSymbol", Ls(toValue(sym.nme)), symName)(k)
case baseSym: BaseTypeSymbol =>
val name = scope.allocateOrGetName(sym)
val (owner, bsym, paramsOpt, auxParams) = (baseSym.defn, defnMap.get(baseSym)) match
case (S(defn), _) => (defn.owner, defn.bsym, defn.paramsOpt, defn.auxParams)
case (_, S(defn: ClsLikeDefn)) => (defn.owner, defn.sym, defn.paramsOpt, defn.auxParams)
case _ =>
raise(ErrorReport(msg"Unable to infer parameters from symbol in staged module, which are necessary to reconstruct class instances: ${sym.toString()}" -> sym.toLoc :: Nil))
return End()

val path: ArgWrappable = pOpt.getOrElse(owner match
case S(owner) => owner.asPath.selSN(sym.nme)
case N => bsym)
baseSym match
case _: ClassSymbol =>
transformParamsOpt(paramsOpt): paramsOpt =>
auxParams.map(ps => transformParamList(ps)).collectApply: auxParams =>
tuple(auxParams): auxParams =>
blockCtor("ClassSymbol", Ls(toValue(name), path, paramsOpt, auxParams), symName)(k)
case _: ModuleOrObjectSymbol =>
blockCtor("ModuleSymbol", Ls(toValue(name), path), symName)(k)
case _ =>
val name = scope.allocateOrGetName(sym)
blockCtor("Symbol", Ls(toValue(name)), symName)(k)
Expand Down Expand Up @@ -156,7 +175,7 @@ class InstrumentationImpl(using State, Raise):
blockCtor("ValueLit", Ls(l), "lit")(k)
case s @ Select(p, Tree.Ident(name)) =>
transformPath(p): x =>
val sym = s.symbol.map(transformSymbol(_))
val sym = s.symbol.map(transformSymbol(_, S(s)))
.getOrElse(blockCtor("Symbol", Ls(toValue(name))))
sym: sym =>
blockCtor("Select", Ls(x, sym), "sel")(k)
Expand Down Expand Up @@ -290,13 +309,13 @@ class InstrumentationImpl(using State, Raise):
transformSymbol(labelSymbol): labelSymbol =>
blockCtor("Break", Ls(labelSymbol))(k(_, ctx))
case _ =>
raise(ErrorReport(msg"Other Blocks not supprted in staged module: ${b.toString()}" -> N :: Nil))
raise(ErrorReport(msg"Other Blocks not supported in staged module: ${b.toString()}" -> N :: Nil))
End()

def transformFunDefn(f: FunDefn)(using Context)(k: Path => Block): Block =
transformBlock(f.body): body =>
if f.params.length != 1 then
raise(WarningReport(msg"Multiple parameter lists are not supported in shape propagation yet." -> f.sym.toLoc :: Nil))
raise(ErrorReport(msg":ftc must be enabled to desugar functions with multiple parameter lists." -> f.sym.toLoc :: Nil))
// maintain parameter names in instrumented code
f.params.map(
_.params.map(p => blockCtor("Symbol", Ls(toValue(p.sym.nme)))).collectApply
Expand All @@ -306,7 +325,7 @@ class InstrumentationImpl(using State, Raise):
tuple(tups): tup =>
blockCtor("FunDefn", Ls(sym, tup, body, toValue(true)))(k)

def applyFunDefn(f: FunDefn): (FunDefn, Block => Block) =
def applyFunDefnInner(f: FunDefn): (FunDefn, Block => Block) =
val genSymName = f.sym.nme + "_instr"
val genSym = BlockMemberSymbol(genSymName, Nil, false)
val sym = f.owner.get.asPath.selSN(genSymName)
Expand All @@ -321,21 +340,17 @@ class InstrumentationImpl(using State, Raise):
val newFun = f.copy(sym = genSym, dSym = dSym, params = Ls(PlainParamList(Nil)), body = newBody)(false)
(newFun, debug)

// TODO: rename as InstrumentationTransformer?
class Instrumentation(using State, Raise) extends BlockTransformer(new SymbolSubst()):
val impl = new InstrumentationImpl

override def applyBlock(b: Block): Block =
super.applyBlock(b) match
// find modules with staged annotation
case Define(c: ClsLikeDefn, rest) if c.companion.exists(_.isym.defn.exists(_.hasStagedModifier.isDefined)) =>
val sym = c.sym.subst
val companion = c.companion.get
val (stagedMethods, debugPrintCode) = companion.methods
.map(impl.applyFunDefn)
.map(applyFunDefnInner)
.unzip
val ctor = FunDefn.withFreshSymbol(S(companion.isym), BlockMemberSymbol("ctor$", Nil), Ls(PlainParamList(Nil)), companion.ctor)(false)
val (stagedCtor, ctorPrint) = impl.applyFunDefn(ctor)
val (stagedCtor, ctorPrint) = applyFunDefnInner(ctor)

val unit = State.runtimeSymbol.asPath.selSN("Unit")
val debugBlock = (ctorPrint :: debugPrintCode).foldRight((Return(unit, true): Block))(_(_))
Expand All @@ -346,7 +361,7 @@ class Instrumentation(using State, Raise) extends BlockTransformer(new SymbolSub
override def applyBlock(b: Block): Block = super.applyBlock(b) match
case Define(c: ClsLikeDefn, rest) if c.companion.isEmpty =>
val (stagedMethods, debugPrintCode) = c.methods
.map(impl.applyFunDefn)
.map(applyFunDefnInner)
.unzip
val newModule = c.copy(methods = c.methods ++ stagedMethods)
Define(newModule, rest)
Expand All @@ -356,3 +371,16 @@ class Instrumentation(using State, Raise) extends BlockTransformer(new SymbolSub
val newModule = c.copy(sym = sym, companion = S(newCompanion))
Define(newModule, rest)
case b => b

def mkDefnMap(b: Block): Unit =
val transformer = new BlockTransformer(new SymbolSubst()):
override def applyDefn(defn: Defn)(k: Defn => Block) = defn match
case c: ClsLikeDefn =>
defnMap.addOne(c.isym, c)
super.applyDefn(defn)(k)
case _ => super.applyDefn(defn)(k)
transformer.applyBlock(b)

def applyBlockFinal(b: Block) =
mkDefnMap(b)
applyBlock(b)
2 changes: 1 addition & 1 deletion hkmc2/shared/src/main/scala/hkmc2/codegen/Lowering.scala
Original file line number Diff line number Diff line change
Expand Up @@ -1088,7 +1088,7 @@ class Lowering()(using Config, TL, Raise, State, Ctx):
if config.funcToCls then Lifter(FirstClassFunctionTransformer().transform(merged)).transform
else merged

val staged = Instrumentation(using summon).applyBlock(funcToCls)
val staged = Instrumentation(using summon).applyBlockFinal(funcToCls)

val res =
if config.tailRecOpt then TailRecOpt().transform(staged)
Expand Down
24 changes: 8 additions & 16 deletions hkmc2/shared/src/test/mlscript-compile/Block.mls
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,11 @@ type Literal = null | undefined | Str | Int | Num | Bool
type ParamList = Array[Symbol]

class Symbol(val name: Str)
type Ident = Symbol
// this is so that we're able to retrieve information about the class from the symbol
class ClassSymbol(val name: Str, val paramsOpt: Opt[ParamList]) extends Symbol(name)
class VirtualClassSymbol(val name: Str) extends Symbol(name)
class ClassSymbol(val name: Str, val value, val paramsOpt: Option[ParamList], val auxParams: Array[ParamList]) extends Symbol(name)
class ModuleSymbol(val name: Str, val value) extends Symbol(name)

// Classes defined in Block.scala

class Arg(val value: Path)

Expand All @@ -37,7 +39,7 @@ class Result with

class Path extends Result with
constructor
Select(val qual: Path, val name: Ident)
Select(val qual: Path, val name: Symbol)
DynSelect(val qual: Path, val fld: Path, val arrayIdx: Bool) // is arrayIdx used?
ValueRef(val l: Symbol)
ValueLit(val lit: Literal)
Expand Down Expand Up @@ -82,22 +84,12 @@ fun showLiteral(l: Literal) =
else l.toString()

fun showSymbol(s: Symbol) =
// console.log("printing " + s)
if s is
ClassSymbol(name, args) then
"ClassSymbol(" + "\"" + name + "\"" +
if args
is Some(args) then ":[" + args.map(showSymbol).join(", ") + "]"
is None then ""
+ ")"
_ then "Symbol(" + "\"" + s.name + "\"" + ")"

fun showIdent(i: Ident) = showSymbol(i)
s.toString()

fun showPath(p: Path): Str =
if p is
Select(qual, name) then
"Select(" + showPath(qual) + ", " + showIdent(name) + ")"
"Select(" + showPath(qual) + ", " + showSymbol(name) + ")"
DynSelect(qual, fld, arrayIdx) then
"DynSelect(" + showPath(qual) + ", " + showPath(fld) + ", " + showBool(arrayIdx) + ")"
ValueRef(l) then
Expand Down
47 changes: 27 additions & 20 deletions hkmc2/shared/src/test/mlscript/block-staging/Functions.mls
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ staged module Expressions with
//│ > FunDefn(Symbol("tup1"), ([]), Return(Tuple([Lit(1), Lit(2)]), false), true)
//│ > FunDefn(Symbol("tup2"), ([]), Return(Tuple([Lit(1), Ref(Symbol("x1"))]), false), true)
//│ > FunDefn(Symbol("dynsel"), ([]), Scoped([Symbol("tmp")], Assign(Symbol("tmp"), Tuple([Lit(1)]), Return(DynSelect(Ref(Symbol("tmp")), Lit(0), false), false))), true)
//│ > FunDefn(Symbol("match1"), ([]), Scoped([Symbol("scrut")], Assign(Symbol("scrut"), Lit(9), Match(Ref(Symbol("scrut")), [Cls(ClassSymbol("Bool"), Select(Ref(Symbol("runtime")), Symbol("unreachable"))) -> Return(Lit(1), false), Lit(8) -> Return(Lit(2), false), Cls(ClassSymbol("Int"), Select(Ref(Symbol("runtime")), Symbol("unreachable"))) -> Return(Lit(3), false)], Return(Lit(0), false), End))), true)
//│ > FunDefn(Symbol("match1"), ([]), Scoped([Symbol("scrut")], Assign(Symbol("scrut"), Lit(9), Match(Ref(Symbol("scrut")), [Cls(VirtualClassSymbol("Bool"), Select(Ref(Symbol("runtime")), Symbol("unreachable"))) -> Return(Lit(1), false), Lit(8) -> Return(Lit(2), false), Cls(VirtualClassSymbol("Int"), Select(Ref(Symbol("runtime")), Symbol("unreachable"))) -> Return(Lit(3), false)], Return(Lit(0), false), End))), true)
//│ > FunDefn(Symbol("match2"), ([]), Scoped([Symbol("a"), Symbol("element1$"), Symbol("element0$")], Match(Ref(Symbol("x1")), [Tup(0) -> Return(Lit(1), false), Tup(2) -> Assign(Symbol("element0$"), Call(Select(Select(Ref(Symbol("runtime")), Symbol("Tuple")), Symbol("get")), [Ref(Symbol("x1")), Lit(0)]), Assign(Symbol("element1$"), Call(Select(Select(Ref(Symbol("runtime")), Symbol("Tuple")), Symbol("get")), [Ref(Symbol("x1")), Lit(1)]), Match(Ref(Symbol("element0$")), [Lit(1) -> Match(Ref(Symbol("element1$")), [Lit(2) -> Return(Lit(2), false)], Assign(Symbol("a"), Ref(Symbol("element0$")), Return(Lit(3), false)), End)], Assign(Symbol("a"), Ref(Symbol("element0$")), Return(Lit(3), false)), End)))], Return(Lit(0), false), End)), true)
//│ x = [1, 2, 3]

Expand All @@ -42,11 +42,11 @@ staged module ClassInstrumentation with
fun inst2() = new NoArg
fun app1() = Outside(1)
fun app2() = Inside(1, 2)
//│ > FunDefn(Symbol("ctor$"), ([]), Define(ClsLikeDefn(ClassSymbol("Inside":[Symbol("a1"), Symbol("b")]), [], TODO), Define(ClsLikeDefn(ClassSymbol("NoArg"), [], TODO), End)), true)
//│ > FunDefn(Symbol("inst1"), ([]), Return(Instantiate(Ref(ClassSymbol("Outside":[Symbol("a")])), [Lit(1)]), false), true)
//│ > FunDefn(Symbol("inst2"), ([]), Return(Instantiate(Select(Ref(Symbol("ClassInstrumentation")), ClassSymbol("NoArg")), []), false), true)
//│ > FunDefn(Symbol("app1"), ([]), Return(Call(Ref(ClassSymbol("Outside":[Symbol("a")])), [Lit(1)]), false), true)
//│ > FunDefn(Symbol("app2"), ([]), Return(Call(Select(Ref(Symbol("ClassInstrumentation")), ClassSymbol("Inside":[Symbol("a1"), Symbol("b")])), [Lit(1), Lit(2)]), false), true)
//│ > FunDefn(Symbol("ctor$"), ([]), Define(ClsLikeDefn(ClassSymbol("Inside", fun Inside { class: class Inside }, Some([Symbol("a1"), Symbol("b")]), []), [], TODO), Define(ClsLikeDefn(ClassSymbol("NoArg", class NoArg, None, []), [], TODO), End)), true)
//│ > FunDefn(Symbol("inst1"), ([]), Return(Instantiate(Ref(ClassSymbol("Outside", fun Outside { class: class Outside }, Some([Symbol("a")]), [])), [Lit(1)]), false), true)
//│ > FunDefn(Symbol("inst2"), ([]), Return(Instantiate(Select(Ref(ModuleSymbol("ClassInstrumentation", class ClassInstrumentation { Inside: fun Inside { class: class Inside }, NoArg: class NoArg })), ClassSymbol("NoArg", class NoArg, None, [])), []), false), true)
//│ > FunDefn(Symbol("app1"), ([]), Return(Call(Ref(ClassSymbol("Outside", fun Outside { class: class Outside }, Some([Symbol("a")]), [])), [Lit(1)]), false), true)
//│ > FunDefn(Symbol("app2"), ([]), Return(Call(Select(Ref(ModuleSymbol("ClassInstrumentation", class ClassInstrumentation { Inside: fun Inside { class: class Inside }, NoArg: class NoArg })), ClassSymbol("Inside", fun Inside { class: class Inside }, Some([Symbol("a1"), Symbol("b")]), [])), [Lit(1), Lit(2)]), false), true)

module Nonstaged with
fun f() = 1
Expand All @@ -60,20 +60,17 @@ staged module CallSubst with
Nonstaged.f()
Staged.f()
//│ > FunDefn(Symbol("ctor$"), ([]), End, true)
//│ > FunDefn(Symbol("call"), ([]), Assign(Symbol("$_no$_symbol$_"), Call(Ref(Symbol("+")), [Lit(1), Lit(1)]), Assign(Symbol("$_no$_symbol$_"), Call(Select(Ref(Symbol("Nonstaged")), Symbol("f")), []), Return(Call(Select(Ref(Symbol("Staged")), Symbol("f_gen")), []), false))), true)
//│ > FunDefn(Symbol("call"), ([]), Assign(Symbol("$_no$_symbol$_"), Call(Ref(Symbol("+")), [Lit(1), Lit(1)]), Assign(Symbol("$_no$_symbol$_"), Call(Select(Ref(ModuleSymbol("Nonstaged", class Nonstaged)), Symbol("f")), []), Return(Call(Select(Ref(ModuleSymbol("Staged", class Staged)), Symbol("f_gen")), []), false))), true)

:w
:ftc
staged module Arguments with
fun f(x) =
x = 1
x
fun g(x)(y, z)() = z
//│ ╔══[WARNING] Multiple parameter lists are not supported in shape propagation yet.
//│ ║ l.70: fun g(x)(y, z)() = z
//│ ╙── ^^^^^^^^^^^^^^^^^^^^
//│ > FunDefn(Symbol("ctor$"), ([]), End, true)
//│ > FunDefn(Symbol("f"), ([Symbol("x")]), Assign(Symbol("x"), Lit(1), Return(Ref(Symbol("x")), false)), true)
//│ > FunDefn(Symbol("g"), ([Symbol("x")],[Symbol("y"), Symbol("z")],[]), Return(Ref(Symbol("z")), false), true)
//│ > FunDefn(Symbol("g"), ([Symbol("x")]), Scoped([Symbol("tmp")], Assign(Symbol("tmp"), Instantiate(Ref(ClassSymbol("Function$", class Function$, None, [])), []), Return(Ref(Symbol("tmp")), false))), true)

staged module OtherBlocks with
fun scope() =
Expand All @@ -87,17 +84,18 @@ staged module OtherBlocks with
3 then 0
else 0
//│ > FunDefn(Symbol("ctor$"), ([]), End, true)
//│ > FunDefn(Symbol("scope"), ([]), Scoped([Symbol("a")], Assign(Symbol("a"), Lit(1), Return(Call(Select(Select(Ref(Symbol("OtherBlocks")), Symbol("scope")), Symbol("locally")), [Ref(Symbol("a"))]), false))), true)
//│ > FunDefn(Symbol("scope"), ([]), Scoped([Symbol("a")], Assign(Symbol("a"), Lit(1), Return(Call(Select(Select(Ref(ModuleSymbol("OtherBlocks", class OtherBlocks)), Symbol("scope")), Symbol("locally")), [Ref(Symbol("a"))]), false))), true)
//│ > FunDefn(Symbol("breakAndLabel"), ([]), Scoped([Symbol("scrut")], Assign(Symbol("scrut"), Lit(1), Match(Ref(Symbol("scrut")), [Lit(2) -> Return(Lit(0), false), Lit(3) -> Return(Lit(0), false)], Return(Lit(0), false), End))), true)

staged module ClassDefs with
class A
//│ > FunDefn(Symbol("ctor$"), ([]), Define(ClsLikeDefn(ClassSymbol("A"), [], TODO), End), true)
//│ > FunDefn(Symbol("ctor$"), ([]), Define(ClsLikeDefn(ClassSymbol("A", class A, None, []), [], TODO), End), true)

staged module ClassFunctions with
class A with
fun f() = 1
//│ > FunDefn(Symbol("ctor$"), ([]), Define(ClsLikeDefn(ClassSymbol("A"), [FunDefn(Symbol("f"), ([]), Return(Lit(1), false), true)], TODO), End), true)
class InnerClass() with
fun f() = 1 + Arguments.f(1)
fun g() = f() // TODO
//│ > FunDefn(Symbol("ctor$"), ([]), Define(ClsLikeDefn(ClassSymbol("InnerClass", fun InnerClass { class: class InnerClass }, Some([]), []), [FunDefn(Symbol("f"), ([]), Scoped([Symbol("tmp")], Assign(Symbol("tmp"), Call(Select(Ref(ModuleSymbol("Arguments", class Arguments)), Symbol("f_gen")), [Lit(1)]), Return(Call(Ref(Symbol("+")), [Lit(1), Ref(Symbol("tmp"))]), false))), true),FunDefn(Symbol("g"), ([]), Return(Call(Select(Ref(ClassSymbol("InnerClass", fun InnerClass { class: class InnerClass }, Some([]), [])), Symbol("f")), []), false), true)], TODO), End), true)

// name collision
class A()
Expand All @@ -113,17 +111,26 @@ module A with
//│ > FunDefn(Symbol("ctor$"), ([]), End, true)
//│ > FunDefn(Symbol("f"), ([]), Return(Lit(1), false), true)

// TODO: force enable :ftc to desugar Label and Break
staged module M with
fun f() = if 1 is
0 then 2
1 then 2
else 2
//│ > FunDefn(Symbol("ctor$"), ([]), End, true)
//│ > FunDefn(Symbol("f"), ([]), Scoped([Symbol("scrut")], Assign(Symbol("scrut"), Lit(1), Match(Ref(Symbol("scrut")), [Lit(0) -> Return(Lit(2), false), Lit(1) -> Return(Lit(2), false)], Return(Lit(2), false), End))), true)

:e
class C(val a)
staged module A with
let x = C(1)
fun f() = set x.a = 0
fun g() = {1 : 2}
//│ ═══[COMPILATION ERROR] Other Blocks not supprted in staged module: AssignField(Ref(term:A/x,None),Ident(a),Lit(IntLit(0)),Return(Select(Ref(tmp:runtime,None),Ident(Unit)),false))
//│ ═══[COMPILATION ERROR] Other Blocks not supported in staged module: AssignField(Ref(term:A/x,None),Ident(a),Lit(IntLit(0)),Return(Select(Ref(tmp:runtime,None),Ident(Unit)),false))
//│ ╔══[COMPILATION ERROR] Other Results not supported in staged module: Record(false,List(RcdArg(Some(Lit(IntLit(1))),Lit(IntLit(2)))))
//│ ║ l.121: fun g() = {1 : 2}
//│ ║ l.128: fun g() = {1 : 2}
//│ ╙── ^
//│ > FunDefn(Symbol("ctor$"), ([]), Assign(Symbol("x"), Call(Ref(ClassSymbol("C":[Symbol("a")])), [Lit(1)]), End), true)
//│ > FunDefn(Symbol("ctor$"), ([]), Assign(Symbol("x"), Call(Ref(ClassSymbol("C", fun C { class: class C }, Some([Symbol("a")]), [])), [Lit(1)]), End), true)
//│ ═══[RUNTIME ERROR] Error: MLscript call unexpectedly returned `undefined`, the forbidden value.

:todo
Expand Down
Loading
Loading