redscript-mc 2.6.2 → 3.0.0

package/src/optimizer/strength_reduction.ts

@@ -0,0 +1,95 @@
+/**
+ * Strength Reduction — MIR algebraic simplification pass.
+ *
+ * Rewrites expensive or redundant arithmetic into cheaper equivalents:
+ * - x * 2   -> x + x
+ * - x * 1   -> x
+ * - x * 0   -> 0
+ * - x + 0   -> x
+ * - x - 0   -> x
+ * - x / 1   -> x
+ * - x ^ 1   -> x
+ * - x * -1  -> -x
+ */
+
+import type { MIRBlock, MIRFunction, MIRInstr, Operand } from '../mir/types'
+
+type PowInstr = { kind: 'pow'; dst: string; a: Operand; b: Operand; sourceLoc?: MIRInstr['sourceLoc'] }
+
+export function strengthReduction(fn: MIRFunction): MIRFunction {
+  return {
+    ...fn,
+    blocks: fn.blocks.map(rewriteBlock),
+  }
+}
+
+function rewriteBlock(block: MIRBlock): MIRBlock {
+  return {
+    ...block,
+    instrs: block.instrs.map(instr => rewriteInstr(instr) ?? instr),
+  }
+}
+
+function rewriteInstr(instr: MIRInstr): MIRInstr | null {
+  const powInstr = asPowInstr(instr)
+  if (powInstr && isConst(powInstr.b, 1))
+    return makeCopy(powInstr.dst, powInstr.a, powInstr.sourceLoc)
+
+  switch (instr.kind) {
+    case 'mul':
+      return rewriteMul(instr)
+    case 'add':
+      if (isConst(instr.a, 0)) return makeCopy(instr.dst, instr.b, instr.sourceLoc)
+      if (isConst(instr.b, 0)) return makeCopy(instr.dst, instr.a, instr.sourceLoc)
+      return null
+    case 'sub':
+      if (isConst(instr.b, 0)) return makeCopy(instr.dst, instr.a, instr.sourceLoc)
+      return null
+    case 'div':
+      if (isConst(instr.b, 1)) return makeCopy(instr.dst, instr.a, instr.sourceLoc)
+      return null
+    default:
+      return null
+  }
+}
+
+function rewriteMul(instr: Extract<MIRInstr, { kind: 'mul' }>): MIRInstr | null {
+  if (isConst(instr.a, 0) || isConst(instr.b, 0))
+    return makeConst(instr.dst, 0, instr.sourceLoc)
+
+  if (isConst(instr.a, 1)) return makeCopy(instr.dst, instr.b, instr.sourceLoc)
+  if (isConst(instr.b, 1)) return makeCopy(instr.dst, instr.a, instr.sourceLoc)
+
+  if (isConst(instr.a, -1)) return makeNeg(instr.dst, instr.b, instr.sourceLoc)
+  if (isConst(instr.b, -1)) return makeNeg(instr.dst, instr.a, instr.sourceLoc)
+
+  if (isConst(instr.a, 2)) return makeAdd(instr.dst, instr.b, instr.b, instr.sourceLoc)
+  if (isConst(instr.b, 2)) return makeAdd(instr.dst, instr.a, instr.a, instr.sourceLoc)
+
+  return null
+}
+
+function isConst(op: Operand, value: number): boolean {
+  return op.kind === 'const' && op.value === value
+}
+
+function asPowInstr(instr: MIRInstr): PowInstr | null {
+  const candidate = instr as MIRInstr | PowInstr
+  return candidate.kind === 'pow' ? candidate : null
+}
+
+function makeCopy(dst: string, src: Operand, sourceLoc?: MIRInstr['sourceLoc']): MIRInstr {
+  return sourceLoc ? { kind: 'copy', dst, src, sourceLoc } : { kind: 'copy', dst, src }
+}
+
+function makeConst(dst: string, value: number, sourceLoc?: MIRInstr['sourceLoc']): MIRInstr {
+  return sourceLoc ? { kind: 'const', dst, value, sourceLoc } : { kind: 'const', dst, value }
+}
+
+function makeNeg(dst: string, src: Operand, sourceLoc?: MIRInstr['sourceLoc']): MIRInstr {
+  return sourceLoc ? { kind: 'neg', dst, src, sourceLoc } : { kind: 'neg', dst, src }
+}
+
+function makeAdd(dst: string, a: Operand, b: Operand, sourceLoc?: MIRInstr['sourceLoc']): MIRInstr {
+  return sourceLoc ? { kind: 'add', dst, a, b, sourceLoc } : { kind: 'add', dst, a, b }
+}

package/src/optimizer/tco.ts

@@ -0,0 +1,267 @@
+/**
+ * Tail Call Optimization (TCO) — MIR optimization pass.
+ *
+ * Detects direct self-tail-calls and converts them into loops, eliminating
+ * the recursive function call entirely.
+ *
+ * ## Definition
+ * A tail call is a `call` instruction whose result is immediately returned
+ * in the same block (i.e., the block's terminator is `return` and the last
+ * non-terminator instruction is `call dst fn args` where fn === current fn,
+ * and the return value is that dst temp).
+ *
+ * ## Transformation
+ * Given a function like:
+ *
+ *   fn factorial(n, acc):
+ *     entry:
+ *       t0 = cmp le n 1
+ *       branch t0 → base_case, recurse
+ *     base_case:
+ *       return acc
+ *     recurse:
+ *       t1 = mul acc n
+ *       t2 = sub n 1
+ *       t3 = call factorial(t2, t1)   ← tail call
+ *       return t3
+ *
+ * We:
+ *   1. Create a new "__tco_entry" preamble block that copies original params
+ *      into "loop parameter" temps (__lp0, __lp1, ...) and jumps to the
+ *      original entry block.
+ *   2. In all blocks (including entry), substitute uses of original param
+ *      temps with loop params, so subsequent iterations use updated values.
+ *   3. In tail-call blocks: evaluate call args (to avoid aliasing), assign
+ *      results to loop params, and jump back to original entry — no call emitted.
+ */
+
+import type {
+  MIRFunction,
+  MIRBlock,
+  MIRInstr,
+  Operand,
+  Temp,
+  BlockId,
+} from '../mir/types'
+
+// ---------------------------------------------------------------------------
+// Public API
+// ---------------------------------------------------------------------------
+
+/**
+ * Attempt to apply TCO to a single function.
+ * Returns the transformed function if any tail calls were optimized,
+ * or the original function unchanged if no self-tail-calls were found.
+ */
+export function tailCallOptimize(fn: MIRFunction): MIRFunction {
+  // Skip macro functions — they use a different calling convention.
+  if (fn.isMacro) return fn
+
+  // Find all blocks that end with a self-tail-call.
+  const tailBlocks = findTailCallBlocks(fn)
+  if (tailBlocks.length === 0) return fn
+
+  return rewriteWithLoop(fn, tailBlocks)
+}
+
+// ---------------------------------------------------------------------------
+// Detection
+// ---------------------------------------------------------------------------
+
+/**
+ * Returns a list of tail-call block descriptors for every block that contains
+ * a self-tail-call: the last non-terminator instruction is a `call` to the
+ * same function, and the block terminator is `return` of that call's result.
+ */
+export function findTailCallBlocks(
+  fn: MIRFunction,
+): Array<{ blockId: BlockId; callIdx: number; argCount: number }> {
+  const results: Array<{ blockId: BlockId; callIdx: number; argCount: number }> = []
+
+  for (const block of fn.blocks) {
+    const term = block.term
+    if (term.kind !== 'return') continue
+
+    // Find the last call instruction in this block
+    let callIdx = -1
+    for (let i = block.instrs.length - 1; i >= 0; i--) {
+      if (block.instrs[i].kind === 'call') {
+        callIdx = i
+        break
+      }
+    }
+    if (callIdx === -1) continue
+
+    const callInstr = block.instrs[callIdx] as Extract<MIRInstr, { kind: 'call' }>
+
+    // Must be a self-call
+    if (callInstr.fn !== fn.name) continue
+
+    // The return value must be exactly the call's dst temp
+    const retVal = term.value
+    if (retVal === null) continue
+    if (retVal.kind !== 'temp') continue
+    if (callInstr.dst === null) continue
+    if (retVal.name !== callInstr.dst) continue
+
+    // Arg count must match param count
+    if (callInstr.args.length !== fn.params.length) continue
+
+    results.push({ blockId: block.id, callIdx, argCount: callInstr.args.length })
+  }
+
+  return results
+}
+
+// ---------------------------------------------------------------------------
+// Transformation
+// ---------------------------------------------------------------------------
+
+function rewriteWithLoop(
+  fn: MIRFunction,
+  tailBlocks: Array<{ blockId: BlockId; callIdx: number; argCount: number }>,
+): MIRFunction {
+  const tailBlockIds = new Set(tailBlocks.map(t => t.blockId))
+
+  // 1. Create loop parameter temps: __lp0, __lp1, ...
+  //    These shadow the original params inside the loop body.
+  const loopParams: Temp[] = fn.params.map((_, i) => `__lp${i}`)
+
+  // 2. Build the entry-preamble block (new entry).
+  //    It copies original params → loop params, then jumps to the old entry.
+  const newEntryBlock: MIRBlock = {
+    id: '__tco_entry',
+    instrs: fn.params.map((p, i) => ({
+      kind: 'copy' as const,
+      dst: loopParams[i],
+      src: { kind: 'temp' as const, name: p.name },
+    })),
+    term: { kind: 'jump', target: fn.entry },
+    preds: [],
+  }
+
+  // 3. Build substitution map: original param → loop param
+  const paramSubst = new Map<Temp, Temp>()
+  fn.params.forEach((p, i) => paramSubst.set(p.name, loopParams[i]))
+
+  // 4. Rewrite each block
+  const newBlocks: MIRBlock[] = []
+
+  for (const block of fn.blocks) {
+    if (tailBlockIds.has(block.id)) {
+      // Rewrite tail-call block: strip call+return, update loop params, jump back
+      const tailInfo = tailBlocks.find(t => t.blockId === block.id)!
+      const callInstr = block.instrs[tailInfo.callIdx] as Extract<MIRInstr, { kind: 'call' }>
+
+      // Instructions before the call (substitute param uses)
+      const preInstrs = block.instrs
+        .slice(0, tailInfo.callIdx)
+        .map(i => substituteInstr(i, paramSubst))
+
+      // Evaluate all args into fresh temps first (avoid aliasing when
+      // e.g. lp0 := lp1; lp1 := lp0 — classic swap problem)
+      const argTemps: Temp[] = callInstr.args.map((_, i) => `__tco_arg${i}`)
+
+      const evalInstrs: MIRInstr[] = callInstr.args.map((arg, i) => ({
+        kind: 'copy' as const,
+        dst: argTemps[i],
+        src: substituteOperand(arg, paramSubst),
+      }))
+
+      const assignInstrs: MIRInstr[] = argTemps.map((argTemp, i) => ({
+        kind: 'copy' as const,
+        dst: loopParams[i],
+        src: { kind: 'temp' as const, name: argTemp },
+      }))
+
+      const newBlock: MIRBlock = {
+        ...block,
+        instrs: [...preInstrs, ...evalInstrs, ...assignInstrs],
+        term: { kind: 'jump', target: fn.entry },
+      }
+      newBlocks.push(newBlock)
+    } else {
+      // Non-tail-call block: substitute param uses with loop params
+      const newInstrs = block.instrs.map(i => substituteInstr(i, paramSubst))
+      const newTerm = substituteInstr(block.term, paramSubst) as MIRBlock['term']
+      newBlocks.push({ ...block, instrs: newInstrs, term: newTerm })
+    }
+  }
+
+  // 5. Recompute predecessors
+  const allBlocks = [newEntryBlock, ...newBlocks]
+  const recomputedBlocks = recomputePreds(allBlocks)
+
+  return {
+    ...fn,
+    entry: '__tco_entry',
+    blocks: recomputedBlocks,
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+function substituteOperand(op: Operand, subst: Map<Temp, Temp>): Operand {
+  if (op.kind === 'temp') {
+    const mapped = subst.get(op.name)
+    if (mapped !== undefined) return { kind: 'temp', name: mapped }
+  }
+  return op
+}
+
+function substituteInstr(instr: MIRInstr, subst: Map<Temp, Temp>): MIRInstr {
+  const s = (op: Operand) => substituteOperand(op, subst)
+
+  switch (instr.kind) {
+    case 'copy': return { ...instr, src: s(instr.src) }
+    case 'add': return { ...instr, a: s(instr.a), b: s(instr.b) }
+    case 'sub': return { ...instr, a: s(instr.a), b: s(instr.b) }
+    case 'mul': return { ...instr, a: s(instr.a), b: s(instr.b) }
+    case 'div': return { ...instr, a: s(instr.a), b: s(instr.b) }
+    case 'mod': return { ...instr, a: s(instr.a), b: s(instr.b) }
+    case 'neg': return { ...instr, src: s(instr.src) }
+    case 'cmp': return { ...instr, a: s(instr.a), b: s(instr.b) }
+    case 'and': return { ...instr, a: s(instr.a), b: s(instr.b) }
+    case 'or': return { ...instr, a: s(instr.a), b: s(instr.b) }
+    case 'not': return { ...instr, src: s(instr.src) }
+    case 'nbt_write': return { ...instr, src: s(instr.src) }
+    case 'nbt_write_dynamic': return { ...instr, indexSrc: s(instr.indexSrc), valueSrc: s(instr.valueSrc) }
+    case 'nbt_read_dynamic': return { ...instr, indexSrc: s(instr.indexSrc) }
+    case 'score_write': return { ...instr, src: s(instr.src) }
+    case 'call': return { ...instr, args: instr.args.map(s) }
+    case 'call_macro': return {
+      ...instr,
+      args: instr.args.map(a => ({ ...a, value: s(a.value) })),
+    }
+    case 'branch': return { ...instr, cond: s(instr.cond) }
+    case 'return': return {
+      ...instr,
+      value: instr.value !== null ? s(instr.value) : null,
+    }
+    default:
+      return instr
+  }
+}
+
+function recomputePreds(blocks: MIRBlock[]): MIRBlock[] {
+  const predsMap = new Map<BlockId, BlockId[]>()
+  for (const b of blocks) predsMap.set(b.id, [])
+
+  for (const b of blocks) {
+    const targets = getTermTargets(b.term)
+    for (const tgt of targets) {
+      const list = predsMap.get(tgt)
+      if (list) list.push(b.id)
+    }
+  }
+
+  return blocks.map(b => ({ ...b, preds: predsMap.get(b.id) ?? [] }))
+}
+
+function getTermTargets(term: MIRInstr): BlockId[] {
+  if (term.kind === 'jump') return [term.target]
+  if (term.kind === 'branch') return [term.then, term.else]
+  return []
+}

package/src/optimizer/unroll.ts

@@ -327,7 +327,7 @@ function substituteInstr(instr: MIRInstr, sub: Map<Temp, Operand>): MIRInstr {
   switch (instr.kind) {
     case 'copy': return { ...instr, src: substituteOp(instr.src, sub) }
     case 'neg': case 'not': return { ...instr, src: substituteOp(instr.src, sub) }
-    case 'add': case 'sub': case 'mul': case 'div': case 'mod':
+    case 'add': case 'sub': case 'mul': case 'div': case 'mod': case 'pow':
     case 'and': case 'or':
       return { ...instr, a: substituteOp(instr.a, sub), b: substituteOp(instr.b, sub) }
     case 'cmp':
@@ -356,7 +356,7 @@ function substituteInstr(instr: MIRInstr, sub: Map<Temp, Operand>): MIRInstr {
 function getInstrDst(instr: MIRInstr): Temp | null {
   switch (instr.kind) {
     case 'const': case 'copy':
-    case 'add': case 'sub': case 'mul': case 'div': case 'mod':
+    case 'add': case 'sub': case 'mul': case 'div': case 'mod': case 'pow':
     case 'neg': case 'cmp': case 'and': case 'or': case 'not':
     case 'nbt_read':
     case 'nbt_read_dynamic':