FazBrowse GitHub Viewer
|
Trending
|
URL:
|
Home
Tools:
[Download Repo ZIP]
[View Raw Code]
[Original HTTPS Page]
modelscript/packages/language/src/codegen/egraph.ts at main · modelscript/modelscript · GitHub
Uh oh!
There was an error while loading.
Please reload this page
.
modelscript
/
modelscript
Public
Uh oh!
There was an error while loading.
Please reload this page
.
Notifications
You must be signed in to change notification settings
Fork
3
Star
12
Code
Issues
1
Pull requests
5
Discussions
Actions
Projects
Security and quality
0
Insights
Additional navigation options
Code
Issues
Pull requests
Discussions
Actions
Projects
Security and quality
Insights
Expand file tree
Breadcrumbs
modelscript
/
packages
/
language
/
src
/
codegen
/
egraph.ts
Copy path
More file actions
More file actions
Latest commit
History
History
History
304 lines (291 loc) · 13.7 KB
Breadcrumbs
modelscript
/
packages
/
language
/
src
/
codegen
/
egraph.ts
Copy path
File metadata and controls
304 lines (291 loc) · 13.7 KB
Raw
Copy raw file
Download raw file
Open symbols panel
Edit and raw actions
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
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
import
{
LanguageOptions
}
from
"../dsl.js"
;
import
{
compileRewriteRules
}
from
"./compile_rules.js"
;
/**
* Generates an AssemblyScript e-graph saturation and Bellman-Ford DP extraction engine.
* Emits zero-GC union-find data structures, dense e-node arrays, hash-consing deduplication,
* and rule matching loops.
*
*
@param
grammar Language configuration options.
*
@param
rules Array of rewrite rule definitions.
*
@returns
AssemblyScript source code string for the e-graph runtime module.
*/
export
function
generateEGraphEngine
(
grammar
:
LanguageOptions
,
rules
:
any
[
]
)
:
string
{
let
out
=
'import { atomicChunkAlloc, getNodeType, getNodeFirstChild, getNodeNextSibling, allocNode } from "./arena";\n'
+
'import { DaeBuilder } from "./dae";\n\n'
+
"// --- EGraph Engine (Zero-GC) ---\n"
+
"export const MAX_ECLASSES: u32 = 65536;\n"
+
"export const MAX_ENODES: u32 = 65536;\n"
+
"export const HASH_CAPACITY: u32 = 65536; // Power of 2\n"
+
"export const HASH_MASK: u32 = HASH_CAPACITY - 1;\n"
+
"export const EMPTY_KEY: u64 = 0xFFFFFFFFFFFFFFFF;\n\n"
+
"export function unwrapNode(node: u32): u32 {\n"
+
" return node;\n"
+
"}\n\n"
+
"// --- Union-Find Disjoint Set ---\n"
+
"let ufParentOffset: u32 = 0;\n"
+
"let ufRankOffset: u32 = 0;\n"
+
"export let ufCount: u32 = 0;\n\n"
+
"export function initEGraph(): void {\n"
+
" if (ufParentOffset == 0) {\n"
+
" ufParentOffset = atomicChunkAlloc(MAX_ECLASSES * 4);\n"
+
" ufRankOffset = atomicChunkAlloc(MAX_ECLASSES);\n"
+
" }\n"
+
" ufCount = 0;\n"
+
"}\n\n"
+
"export function ufMakeSet(): u32 {\n"
+
" if (ufCount >= MAX_ECLASSES) return 0xFFFFFFFF;\n"
+
" let id = ufCount++;\n"
+
" store<u32>(ufParentOffset + id * 4, id);\n"
+
" store<u8>(ufRankOffset + id, 0);\n"
+
" return id;\n"
+
"}\n\n"
+
"export function ufFind(x: u32): u32 {\n"
+
" if (x >= MAX_ECLASSES || x == 0xFFFFFFFF || ufParentOffset == 0) return x;\n"
+
" let root = x;\n"
+
" while (true) {\n"
+
" let parent = load<u32>(ufParentOffset + root * 4);\n"
+
" if (parent == root || parent >= MAX_ECLASSES) break;\n"
+
" root = parent;\n"
+
" }\n"
+
" let curr = x;\n"
+
" while (curr != root && curr < MAX_ECLASSES) {\n"
+
" let nxt = load<u32>(ufParentOffset + curr * 4);\n"
+
" store<u32>(ufParentOffset + curr * 4, root);\n"
+
" curr = nxt;\n"
+
" }\n"
+
" return root;\n"
+
"}\n\n"
+
"export function ufUnion(a: u32, b: u32): u32 {\n"
+
" let rootA = ufFind(a);\n"
+
" let rootB = ufFind(b);\n"
+
" if (rootA == rootB) return rootA;\n"
+
" let rankA = load<u8>(ufRankOffset + rootA);\n"
+
" let rankB = load<u8>(ufRankOffset + rootB);\n"
+
" if (rankA < rankB) {\n"
+
" store<u32>(ufParentOffset + rootA * 4, rootB);\n"
+
" return rootB;\n"
+
" } else if (rankA > rankB) {\n"
+
" store<u32>(ufParentOffset + rootB * 4, rootA);\n"
+
" return rootA;\n"
+
" } else {\n"
+
" store<u32>(ufParentOffset + rootB * 4, rootA);\n"
+
" store<u8>(ufRankOffset + rootA, rankA + 1);\n"
+
" return rootA;\n"
+
" }\n"
+
"}\n\n"
+
"// --- Dense E-Node Storage & Hash-Consing Deduplication Table ---\n"
+
"export let eNodeKeysOffset: u32 = 0;\n"
+
"export let eNodeClassesOffset: u32 = 0;\n"
+
"export let eNodeCount: u32 = 0;\n\n"
+
"let hashKeysOffset: u32 = 0;\n"
+
"let hashValsOffset: u32 = 0;\n"
+
"export let hashOccupied: u32 = 0;\n\n"
+
"export function initHashCons(): void {\n"
+
" if (hashKeysOffset == 0) {\n"
+
" hashKeysOffset = atomicChunkAlloc(HASH_CAPACITY * 8);\n"
+
" hashValsOffset = atomicChunkAlloc(HASH_CAPACITY * 4);\n"
+
" eNodeKeysOffset = atomicChunkAlloc(MAX_ENODES * 8);\n"
+
" eNodeClassesOffset = atomicChunkAlloc(MAX_ENODES * 4);\n"
+
" }\n"
+
" eNodeCount = 0;\n"
+
" hashOccupied = 0;\n"
+
" // Set all key slots to EMPTY_KEY sentinel (0xFFFFFFFFFFFFFFFF)\n"
+
" memory.fill(hashKeysOffset, 0xFF, HASH_CAPACITY * 8);\n"
+
"}\n\n"
+
"export function hashProbe(key: u64): u32 {\n"
+
" let h = (key ^ (key >> 32)) as u32;\n"
+
" h = ((h >> 16) ^ h) * 0x45d9f3b;\n"
+
" h = ((h >> 16) ^ h);\n"
+
" return h & HASH_MASK;\n"
+
"}\n\n"
+
"export function hashFind(key: u64): u32 {\n"
+
" let slot = hashProbe(key);\n"
+
" let guard: u32 = 0;\n"
+
" while (guard < HASH_CAPACITY) {\n"
+
" let storedKey = load<u64>(hashKeysOffset + slot * 8);\n"
+
" if (storedKey == EMPTY_KEY) return 0xFFFFFFFF; // Empty slot\n"
+
" if (storedKey == key) return load<u32>(hashValsOffset + slot * 4);\n"
+
" slot = (slot + 1) & HASH_MASK;\n"
+
" guard++;\n"
+
" }\n"
+
" return 0xFFFFFFFF;\n"
+
"}\n\n"
+
"export function hashInsert(key: u64, val: u32): void {\n"
+
" let slot = hashProbe(key);\n"
+
" let guard: u32 = 0;\n"
+
" while (guard < HASH_CAPACITY) {\n"
+
" let storedKey = load<u64>(hashKeysOffset + slot * 8);\n"
+
" if (storedKey == EMPTY_KEY) {\n"
+
" store<u64>(hashKeysOffset + slot * 8, key);\n"
+
" store<u32>(hashValsOffset + slot * 4, val);\n"
+
" if (eNodeCount < MAX_ENODES) {\n"
+
" store<u64>(eNodeKeysOffset + eNodeCount * 8, key);\n"
+
" store<u32>(eNodeClassesOffset + eNodeCount * 4, val);\n"
+
" eNodeCount++;\n"
+
" }\n"
+
" hashOccupied++;\n"
+
" return;\n"
+
" }\n"
+
" if (storedKey == key) {\n"
+
" store<u32>(hashValsOffset + slot * 4, val);\n"
+
" return;\n"
+
" }\n"
+
" slot = (slot + 1) & HASH_MASK;\n"
+
" guard++;\n"
+
" }\n"
+
"}\n\n"
+
"export function rebuildEGraph(): void {\n"
+
" memory.fill(hashKeysOffset, 0xFF, HASH_CAPACITY * 8);\n"
+
" hashOccupied = 0;\n"
+
" let writeIdx: u32 = 0;\n"
+
" for (let i: u32 = 0; i < eNodeCount; i++) {\n"
+
" let key = load<u64>(eNodeKeysOffset + i * 8);\n"
+
" let eClass = ufFind(load<u32>(eNodeClassesOffset + i * 4));\n"
+
" let op = (key >> 48) as u16;\n"
+
" let left = ((key >> 24) & 0xFFFFFF) as u32;\n"
+
" let right = (key & 0xFFFFFF) as u32;\n\n"
+
" if ((op >= 1280 && op <= 1284) || (op >= 1536 && op <= 1539) || (op >= 1792 && op <= 1793)) {\n"
+
" left = ufFind(left);\n"
+
" right = ufFind(right);\n"
+
" key = ((op as u64) << 48) | (((left & 0xFFFFFF) as u64) << 24) | ((right & 0xFFFFFF) as u64);\n"
+
" } else if ((op >= 1024 && op <= 1027) || (op >= 1800 && op <= 1810)) {\n"
+
" left = ufFind(left);\n"
+
" key = ((op as u64) << 48) | (((left & 0xFFFFFF) as u64) << 24);\n"
+
" }\n\n"
+
" let existing = hashFind(key);\n"
+
" if (existing != 0xFFFFFFFF) {\n"
+
" ufUnion(eClass, existing);\n"
+
" } else {\n"
+
" let slot = hashProbe(key);\n"
+
" let guard: u32 = 0;\n"
+
" while (guard < HASH_CAPACITY) {\n"
+
" let storedKey = load<u64>(hashKeysOffset + slot * 8);\n"
+
" if (storedKey == EMPTY_KEY) {\n"
+
" store<u64>(hashKeysOffset + slot * 8, key);\n"
+
" store<u32>(hashValsOffset + slot * 4, eClass);\n"
+
" hashOccupied++;\n"
+
" break;\n"
+
" }\n"
+
" slot = (slot + 1) & HASH_MASK;\n"
+
" guard++;\n"
+
" }\n"
+
" store<u64>(eNodeKeysOffset + writeIdx * 8, key);\n"
+
" store<u32>(eNodeClassesOffset + writeIdx * 4, eClass);\n"
+
" writeIdx++;\n"
+
" }\n"
+
" }\n"
+
" eNodeCount = writeIdx;\n"
+
"}\n\n"
+
"export function isConstant(eClass: u32, val: f64): boolean {\n"
+
" let root = ufFind(eClass);\n"
+
" let floatBits = reinterpret<u64>(val);\n"
+
" let keyReal: u64 = ((512 as u64) << 48) | (floatBits >>> 16);\n"
+
" let classReal = hashFind(keyReal);\n"
+
" if (classReal != 0xFFFFFFFF && ufFind(classReal) == root) return true;\n"
+
" let keyInt: u64 = ((256 as u64) << 48) | ((val as u32) & 0xFFFFFFFF);\n"
+
" let classInt = hashFind(keyInt);\n"
+
" if (classInt != 0xFFFFFFFF && ufFind(classInt) == root) return true;\n"
+
" return false;\n"
+
"}\n\n"
;
out
+=
"export function addENode(exprId: u32, dae: DaeBuilder): u32 {\n"
;
out
+=
" if (exprId == 0xFFFFFFFF) return 0xFFFFFFFF;\n"
;
out
+=
" let exprOffset = exprId * 4;\n"
;
out
+=
" let kind = dae.exprData.get(exprOffset + 0);\n"
;
out
+=
" let data1 = dae.exprData.get(exprOffset + 1);\n"
;
out
+=
" let data2 = dae.exprData.get(exprOffset + 2);\n\n"
;
out
+=
" if (kind == 0) {\n"
;
// Name
out
+=
" let key: u64 = (data1 as u64);\n"
;
out
+=
" let existing = hashFind(key);\n"
;
out
+=
" if (existing != 0xFFFFFFFF) return ufFind(existing);\n"
;
out
+=
" let id = ufMakeSet();\n"
;
out
+=
" hashInsert(key, id);\n"
;
out
+=
" return id;\n"
;
out
+=
" }\n"
;
out
+=
" if (kind == 1) {\n"
;
// IntLiteral
out
+=
" let key: u64 = ((256 as u64) << 48) | ((data1 as u32) as u64);\n"
;
out
+=
" let existing = hashFind(key);\n"
;
out
+=
" if (existing != 0xFFFFFFFF) return ufFind(existing);\n"
;
out
+=
" let id = ufMakeSet();\n"
;
out
+=
" hashInsert(key, id);\n"
;
out
+=
" return id;\n"
;
out
+=
" }\n"
;
out
+=
" if (kind == 2) {\n"
;
// RealLiteral
out
+=
" let lo = data1 as u64;\n"
;
out
+=
" let hi = data2 as u64;\n"
;
out
+=
" let floatBits: u64 = lo | (hi << 32);\n"
;
out
+=
" let key: u64 = ((512 as u64) << 48) | (floatBits >>> 16);\n"
;
out
+=
" let existing = hashFind(key);\n"
;
out
+=
" if (existing != 0xFFFFFFFF) return ufFind(existing);\n"
;
out
+=
" let id = ufMakeSet();\n"
;
out
+=
" hashInsert(key, id);\n"
;
out
+=
" return id;\n"
;
out
+=
" }\n"
;
out
+=
" if (kind == 3) {\n"
;
// BoolLiteral
out
+=
" let key: u64 = ((768 as u64) << 48) | ((data1 != 0 ? 1 : 0) as u64);\n"
;
out
+=
" let existing = hashFind(key);\n"
;
out
+=
" if (existing != 0xFFFFFFFF) return ufFind(existing);\n"
;
out
+=
" let id = ufMakeSet();\n"
;
out
+=
" hashInsert(key, id);\n"
;
out
+=
" return id;\n"
;
out
+=
" }\n"
;
out
+=
" if (kind == 5) {\n"
;
// Binary
out
+=
" let leftId = dae.exprData.get(exprOffset + 2);\n"
;
out
+=
" let rightId = dae.exprData.get(exprOffset + 3);\n"
;
out
+=
" let leftClass = addENode(leftId, dae);\n"
;
out
+=
" let rightClass = addENode(rightId, dae);\n"
;
out
+=
" let opType = (kind << 8) | data1;\n"
;
out
+=
" let key: u64 = ((opType as u64) << 48) | (((ufFind(leftClass) & 0xFFFFFF) as u64) << 24) | ((ufFind(rightClass) & 0xFFFFFF) as u64);\n"
;
out
+=
" let existing = hashFind(key);\n"
;
out
+=
" if (existing != 0xFFFFFFFF) return ufFind(existing);\n"
;
out
+=
" let id = ufMakeSet();\n"
;
out
+=
" hashInsert(key, id);\n"
;
out
+=
" return id;\n"
;
out
+=
" }\n"
;
out
+=
" if (kind == 6) {\n"
;
// Unary
out
+=
" let childId = dae.exprData.get(exprOffset + 2);\n"
;
out
+=
" let childClass = addENode(childId, dae);\n"
;
out
+=
" let opType = (kind << 8) | data1;\n"
;
out
+=
" let key: u64 = ((opType as u64) << 48) | (((ufFind(childClass) & 0xFFFFFF) as u64) << 24);\n"
;
out
+=
" let existing = hashFind(key);\n"
;
out
+=
" if (existing != 0xFFFFFFFF) return ufFind(existing);\n"
;
out
+=
" let id = ufMakeSet();\n"
;
out
+=
" hashInsert(key, id);\n"
;
out
+=
" return id;\n"
;
out
+=
" }\n"
;
out
+=
" if (kind == 7) {\n"
;
// Call
out
+=
" let childId = dae.exprData.get(exprOffset + 2);\n"
;
out
+=
" let childClass = addENode(childId, dae);\n"
;
out
+=
" let opType = 1800 + (data1 as u16);\n"
;
out
+=
" let key: u64 = ((opType as u64) << 48) | (((ufFind(childClass) & 0xFFFFFF) as u64) << 24);\n"
;
out
+=
" let existing = hashFind(key);\n"
;
out
+=
" if (existing != 0xFFFFFFFF) return ufFind(existing);\n"
;
out
+=
" let id = ufMakeSet();\n"
;
out
+=
" hashInsert(key, id);\n"
;
out
+=
" return id;\n"
;
out
+=
" }\n"
;
out
+=
" return 0xFFFFFFFF;\n"
;
out
+=
"}\n"
;
if
(
rules
&&
rules
.
length
>
0
)
{
out
+=
compileRewriteRules
(
rules
)
;
}
else
{
out
+=
"export function saturateEGraph(): void {}\n"
;
out
+=
"export function initDPExtractor(): void {}\n"
;
out
+=
"export function extractAst(rootClass: u32, dae: DaeBuilder): u32 { return 0; }\n"
;
}
out
+=
"\n// --- Global AST Simplification ---\n"
+
"export function simplifyAst(exprId: u32, dae: DaeBuilder): u32 {\n"
+
" initEGraph();\n"
+
" initHashCons();\n"
+
" let rootClass = addENode(exprId, dae);\n"
+
" if (rootClass == 0xFFFFFFFF) return exprId;\n"
+
" saturateEGraph();\n"
+
" initDPExtractor();\n"
+
" let simplifiedAst = extractAst(rootClass, dae);\n"
+
" if (simplifiedAst == 0xFFFFFFFF) return exprId;\n"
+
" return simplifiedAst;\n"
+
"}\n"
;
return
out
;
}
Back
|
FazBrowse Home
|
New Git URL