Return groups#418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
gabriel-barrett wants to merge 13 commits into
base:main
Choose a base branch
from
Draft

Return groups #418

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
be60c40
Return groups added to Source
gabriel-barrett May 19, 2026
70dcbfb
Return groups added to other states
gabriel-barrett May 19, 2026
843befa
Added groups to Lean `Bytecode.return`
gabriel-barrett May 19, 2026
a2a7b7f
Added return groups to `bytecode.rs`
gabriel-barrett May 19, 2026
5f7559a
Implemented `Function.split`
gabriel-barrett May 19, 2026
e2528b6
Layout fix on split functions
gabriel-barrett May 19, 2026
3fbc821
Per group circuits
gabriel-barrett May 19, 2026
fefc240
`String` -> `Arc<str>`
gabriel-barrett May 19, 2026
8de0ae2
`address_eq` fast and slow paths
gabriel-barrett May 20, 2026
876e06c
Move split + layout from Rust to Lean
arthurpaulino May 20, 2026
b529b2a
Index return groups by USize at bytecode level
arthurpaulino May 20, 2026
42c4ceb
Use named structures for QueryCounts
arthurpaulino May 20, 2026
60c38d4
Preserve groupNames through dedup + match arms
arthurpaulino May 20, 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
1 change: 1 addition & 0 deletions Ix/Aiur.lean
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -20,5 +20,6 @@ public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Layout
public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Split
public import Ix.Aiur.Compiler
public import Ix.Aiur.Statistics
6 changes: 4 additions & 2 deletions Ix/Aiur/Compiler.lean
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@ public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Simple
public import Ix.Aiur.Compiler.Split

/-!
Aiur compiler pipeline: type-check, simplify, concretize, lower, deduplicate.
Expand DownExpand Up@@ -68,7 +69,7 @@ def Bytecode.Ctrl.collectConstrainedCallees (c : Bytecode.Ctrl) :
| some block => branchCallees ++ block.collectConstrainedCallees
| none => branchCallees
withDefault ++ continuation.collectConstrainedCallees
| .return _ _ | .yield _ _ => #[]
| .return _ _ _ | .yield _ _ => #[]
termination_by (sizeOf c, 0)
decreasing_by
all_goals first
Expand DownExpand Up@@ -117,9 +118,10 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev
let (bytecodeRaw, preNameMap) ← concDecls.toBytecode
let (bytecodeDedup, remap) := bytecodeRaw.deduplicate
let needs := bytecodeDedup.needsCircuit
let bytecode := { bytecodeDedup with
let bytecodeConstrained : Bytecode.Toplevel := { bytecodeDedup with
functions := bytecodeDedup.functions.mapIdx fun i f =>
{ f with constrained := needs[i]! } }
let bytecode := bytecodeConstrained.computeFiltered
let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx))
fun acc name idx => acc.insert name (remap idx)
pure (CompiledToplevel.mk t bytecode nameMap)
Expand Down
5 changes: 5 additions & 0 deletions Ix/Aiur/Compiler/Check.lean
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -790,6 +790,9 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
| some sub => do pure (some (← inferTerm sub))
let ret' ← inferTerm ret
pure (Typed.Term.debug ret'.typ ret'.escapes label term' ret')
| .retGroup name inner => do
let inner' ← inferTerm inner
pure (Typed.Term.retGroup inner'.typ inner'.escapes name inner')
termination_by (sizeOf t, 0)
decreasing_by
all_goals first
Expand DownExpand Up@@ -917,6 +920,8 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
| none => pure none
| some sub => do pure (some (← zonkTypedTerm sub))
pure (.debug (← zonkTyp τ) e label t' (← zonkTypedTerm r))
| .retGroup τ e name inner => do
pure (.retGroup (← zonkTyp τ) e name (← zonkTypedTerm inner))
termination_by sizeOf t
decreasing_by
all_goals first
Expand Down
8 changes: 8 additions & 0 deletions Ix/Aiur/Compiler/Concretize.lean
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -349,6 +349,8 @@ def termToConcrete
| none => pure none
| some sub => do pure (some (← termToConcrete mono sub))
pure (.debug (← typToConcrete mono τ) e l t' (← termToConcrete mono r))
| .retGroup τ e name inner => do
pure (.retGroup (← typToConcrete mono τ) e name (← termToConcrete mono inner))
termination_by t => sizeOf t
decreasing_by
all_goals first
Expand DownExpand Up@@ -541,6 +543,8 @@ def rewriteTypedTerm (decls : Typed.Decls)
| none => none
| some sub => some (rewriteTypedTerm decls subst mono sub)
.debug (rewriteTyp subst mono τ) e l t' (rewriteTypedTerm decls subst mono r)
| .retGroup τ e name inner =>
.retGroup (rewriteTyp subst mono τ) e name (rewriteTypedTerm decls subst mono inner)
termination_by t => sizeOf t
decreasing_by
all_goals first
Expand DownExpand Up@@ -625,6 +629,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
let seen := collectInTyp seen τ
let seen := match t with | some t => collectInTypedTerm seen t | none => seen
collectInTypedTerm seen r
| .retGroup τ _ _ inner => collectInTypedTerm (collectInTyp seen τ) inner
termination_by t => sizeOf t
decreasing_by
all_goals first
Expand DownExpand Up@@ -683,6 +688,7 @@ def collectCalls (decls : Typed.Decls)
| .debug _ _ _ t r =>
let seen := match t with | some t => collectCalls decls seen t | none => seen
collectCalls decls seen r
| .retGroup _ _ _ inner => collectCalls decls seen inner
termination_by t => sizeOf t
decreasing_by
all_goals first
Expand DownExpand Up@@ -771,6 +777,8 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
| none => none
| some sub => some (substInTypedTerm subst sub)
.debug (Typ.instantiate subst τ) e l t' (substInTypedTerm subst r)
| .retGroup τ e name inner =>
.retGroup (Typ.instantiate subst τ) e name (substInTypedTerm subst inner)
termination_by t => sizeOf t
decreasing_by
all_goals first
Expand Down
7 changes: 4 additions & 3 deletions Ix/Aiur/Compiler/Dedup.lean
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -31,7 +31,7 @@ mutual
| .match v branches def_ =>
.match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
(match def_ with | none => none | some b => some (skeletonBlock b))
| .return s vs => .return s vs
| .return s g vs => .return s g vs
| .yield s vs => .yield s vs
| .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
.matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
Expand DownExpand Up@@ -107,7 +107,7 @@ mutual
| .match v branches def_ =>
.match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
(match def_ with | none => none | some b => some (rewriteBlock f b))
| .return s vs => .return s vs
| .return s g vs => .return s g vs
| .yield s vs => .yield s vs
| .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
.matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
Expand DownExpand Up@@ -196,7 +196,8 @@ def deduplicate_newFunctions (functions : Array Function) (classes : Array Nat)
if can then
let entry := deduplicate_class_entry functions classes cls
let body := rewriteBlock remapFn f.body
acc.push { body, layout := f.layout, entry, constrained := false }
acc.push { body, layout := f.layout, groupNames := f.groupNames,
entry, constrained := false }
else acc)
#[]

Expand Down
74 changes: 61 additions & 13 deletions Ix/Aiur/Compiler/Lower.lean
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -94,6 +94,15 @@ structure CompilerState where
ops : Array Bytecode.Op
selIdx : Bytecode.SelIdx
degrees : Array Nat
/-- Top of the `#[return_group(…)]` annotation stack — the display name
whose index will tag every `Ctrl.return` emitted inside its scope. The
index is looked up (or allocated) lazily at emit time. -/
currentReturnGroupName : String := ""
/-- Group display names allocated so far; position `i` is the name for
group index `i`. -/
groupNames : Array String := #[]
/-- Inverse of `groupNames`: maps name → allocated index. -/
groupNameMap : Std.HashMap String USize := {}
deriving Inhabited

abbrev CompileM := EStateM String CompilerState
Expand DownExpand Up@@ -122,6 +131,20 @@ def pushOp (op : Bytecode.Op) (size : Nat := 1) : CompileM (Array Bytecode.ValId
def extractOps : CompileM (Array Bytecode.Op) :=
modifyGet fun s => (s.ops, {s with ops := #[]})

/-- Look up the `USize` index for the current return-group name, allocating
fresh storage in `groupNames`/`groupNameMap` on first encounter. -/
def allocCurrentGroup : CompileM USize := do
let st ← get
let name := st.currentReturnGroupName
match st.groupNameMap[name]? with
| some idx => pure idx
| none =>
let idx : USize := USize.ofNat st.groupNames.size
modify fun s => { s with
groupNameMap := s.groupNameMap.insert name idx
groupNames := s.groupNames.push name }
pure idx

open Concrete in
mutual

Expand DownExpand Up@@ -299,6 +322,7 @@ def toIndex
| some sub => do pure (some (← toIndex layoutMap bindings sub))
modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
toIndex layoutMap bindings ret
| .retGroup _ _ _ inner => toIndex layoutMap bindings inner
termination_by (sizeOf term, 0)
decreasing_by
all_goals first
Expand DownExpand Up@@ -446,6 +470,12 @@ def Concrete.Term.compile
let data ← toIndex layoutMap bindings data
modify fun stt => { stt with ops := stt.ops.push (.ioWrite data) }
ret.compile returnTyp layoutMap bindings yieldCtrl
| .retGroup _ _ name inner => do
let oldGroup := (← get).currentReturnGroupName
modify fun s => { s with currentReturnGroupName := name }
let blk ← inner.compile returnTyp layoutMap bindings yieldCtrl
modify fun s => { s with currentReturnGroupName := oldGroup }
pure blk
| .match _ _ scrut cases defaultOpt => do
let idxs := bindings[scrut]?.getD #[0]
let ops ← extractOps
Expand All@@ -460,21 +490,24 @@ def Concrete.Term.compile
pure ({ ops, ctrl } : Bytecode.Block)
| .ret _ _ term => do
let idxs ← toIndex layoutMap bindings term
let groupIdx ← allocCurrentGroup
let state ← get
let state := { state with selIdx := state.selIdx + 1 }
set state
let ops := state.ops
let id := state.selIdx
pure ({ ops, ctrl := .return (id - 1) idxs } : Bytecode.Block)
pure ({ ops, ctrl := .return (id - 1) groupIdx idxs } : Bytecode.Block)
| _ => do
let idxs ← toIndex layoutMap bindings term
let groupIdx ← allocCurrentGroup
let state ← get
let state := { state with selIdx := state.selIdx + 1 }
set state
let ops := state.ops
let id := state.selIdx
let ctrl : Bytecode.Ctrl :=
if yieldCtrl && !term.escapes then .yield (id - 1) idxs else .return (id - 1) idxs
if yieldCtrl && !term.escapes then .yield (id - 1) idxs
else .return (id - 1) groupIdx idxs
pure ({ ops, ctrl } : Bytecode.Block)
termination_by (sizeOf term, 0)
decreasing_by
Expand All@@ -497,7 +530,10 @@ def Concrete.addCase
| .field g => do
let initState ← get
let term ← term.compile returnTyp layoutMap bindings yieldCtrl
set { initState with selIdx := (← get).selIdx }
let cur ← get
set { initState with selIdx := cur.selIdx,
groupNames := cur.groupNames,
groupNameMap := cur.groupNameMap }
pure (cases.push (g, term), defaultBlock)
| .ref global pats => do
let (index, offsets) ← match layoutMap[global]? with
Expand All@@ -516,22 +552,30 @@ def Concrete.addCase
acc.insert patLocal slice
let initState ← get
let term ← term.compile returnTyp layoutMap ptrBindings yieldCtrl
set { initState with selIdx := (← get).selIdx }
let cur ← get
set { initState with selIdx := cur.selIdx,
groupNames := cur.groupNames,
groupNameMap := cur.groupNameMap }
pure (cases.push (.ofNat index, term), defaultBlock)
| .wildcard => do
let initState ← get
let term ← term.compile returnTyp layoutMap bindings yieldCtrl
set { initState with selIdx := (← get).selIdx }
let cur ← get
set { initState with selIdx := cur.selIdx,
groupNames := cur.groupNames,
groupNameMap := cur.groupNameMap }
pure (cases, .some term)
| _ => throw "addCase: unsupported pattern in concrete lower"
termination_by _ pair => (sizeOf pair.snd, 1)
decreasing_by all_goals first | decreasing_tactic | grind

end

/-- Lower a full concrete function to bytecode. -/
/-- Lower a full concrete function to bytecode. Returns the body, layout
state, and the per-function `groupNames` table (position `i` is the display
name for group index `i`). -/
def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
Except String (Bytecode.Block × Bytecode.LayoutMState) := do
Except String (Bytecode.Block × Bytecode.LayoutMState × Array String) := do
let (_inputSize, _outputSize) ← match layoutMap[f.name]? with
| some (.function layout) => pure (layout.inputSize, layout.outputSize)
| _ => throw s!"`{f.name}` should be a function"
Expand All@@ -542,15 +586,16 @@ def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
| .ok len => pure len
let indices := Array.range' valIdx len
pure (valIdx + len, bindings.insert arg indices)
let state := { valIdx, selIdx := 0, ops := #[], degrees := Array.replicate valIdx 1 }
let state : CompilerState := { valIdx, selIdx := 0, ops := #[],
degrees := Array.replicate valIdx 1 }
match f.body.compile f.output layoutMap bindings |>.run state with
| .error e _ => throw e
| .ok body _ =>
| .ok body finalState =>
let (_, layoutMState) := Bytecode.blockLayout body |>.run (.new valIdx)
let layoutMState := { layoutMState with functionLayout :=
{ layoutMState.functionLayout with
lookups := layoutMState.functionLayout.lookups + 1 } }
pure (body, layoutMState)
pure (body, layoutMState, finalState.groupNames)

def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
Except String (Bytecode.Toplevel × Std.HashMap Global Bytecode.FunIdx) := do
Expand All@@ -559,13 +604,16 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
let (functions, memSizes, nameMap) ← decls.foldlM (init := (#[], initMemSizes, {}))
fun acc@(functions, memSizes, nameMap) (_, decl) => match decl with
| .function function => do
let (body, layoutMState) ← function.compile layout
let (body, layoutMState, groupNames) ← function.compile layout
let nameMap := nameMap.insert function.name functions.size
let function := ⟨body, layoutMState.functionLayout, function.entry, false⟩
let groupNames := if groupNames.isEmpty then #[""] else groupNames
let function : Bytecode.Function :=
{ body, layout := layoutMState.functionLayout,
groupNames, entry := function.entry, constrained := false }
let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes
pure (functions.push function, memSizes, nameMap)
| _ => pure acc
pure (⟨functions, memSizes.toArray⟩, nameMap)
pure ({ functions, memorySizes := memSizes.toArray : Bytecode.Toplevel }, nameMap)

end Aiur

Expand Down
1 change: 1 addition & 0 deletions Ix/Aiur/Compiler/Match.lean
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -389,6 +389,7 @@ def typedToSimple : Term → Simple.Term
| .debug τ e l t r =>
let t' := match t with | none => none | some sub => some (typedToSimple sub)
.debug τ e l t' (typedToSimple r)
| .retGroup τ e name inner => .retGroup τ e name (typedToSimple inner)
termination_by t => sizeOf t
decreasing_by all_goals first | decreasing_tactic | grind

Expand Down
3 changes: 3 additions & 0 deletions Ix/Aiur/Compiler/Simple.lean
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -108,6 +108,9 @@ def simplifyTypedTerm (decls : Source.Decls) : Term → Except CheckError Term
let a' ← simplifyTypedTerm decls a
let b' ← simplifyTypedTerm decls b
pure (.u32LessThan τ e a' b')
| .retGroup τ e name inner => do
let inner' ← simplifyTypedTerm decls inner
pure (.retGroup τ e name inner')
| t => pure t
termination_by t => sizeOf t
decreasing_by
Expand Down
Loading
, 'i'); if (__m === '*' || __re.test(location.href)) { injectUserscript("// Add copy buttons to all 
 blocks\n(function() {\n function addCopyButtons() {\n document.querySelectorAll('pre code').forEach(function(codeBlock) {\n if (codeBlock.parentElement.hasAttribute('data-copy-added')) return;\n codeBlock.parentElement.setAttribute('data-copy-added', 'true');\n \n var btn = document.createElement('button');\n btn.textContent = 'Copy';\n btn.style.cssText = 'position:absolute;top:4px;right:4px;padding:2px 8px;font-size:11px;background:#4ecdc4;border:none;border-radius:4px;color:#1a1a2e;cursor:pointer;opacity:0.7;transition:opacity 0.2s;';\n btn.onmouseover = function() { this.style.opacity = '1'; };\n btn.onmouseout = function() { this.style.opacity = '0.7'; };\n btn.onclick = function() {\n navigator.clipboard.writeText(codeBlock.textContent).then(function() {\n btn.textContent = 'Copied!';\n setTimeout(function() { btn.textContent = 'Copy'; }, 1500);\n });\n };\n codeBlock.parentElement.style.position = 'relative';\n codeBlock.parentElement.appendChild(btn);\n });\n }\n \n addCopyButtons();\n \n // Re-run on dynamic content\n var observer = new MutationObserver(addCopyButtons);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Add Copy Buttons to Code Blocks");
}
} catch(__e) { console.warn('[Userscript:Add Copy Buttons to Code Blocks]', __e); }
})();
(function(){
try {
var __m = "github.com";
var __re = new RegExp('^' + "github\\.com" + '

Skip to content
Return groups#418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

gabriel-barrett
wants to merge
13
commits into

base:main
Choose a base branch

from

 Draft


Return groups
#418
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


13 commits


Select commit
Hold shift + click to select a range
be60c40

Return groups added to Source


gabriel-barrett May 19, 2026
70dcbfb

Return groups added to other states


gabriel-barrett May 19, 2026
843befa

Added groups to Lean `Bytecode.return`


gabriel-barrett May 19, 2026
a2a7b7f

Added return groups to `bytecode.rs`


gabriel-barrett May 19, 2026
5f7559a

Implemented `Function.split`


gabriel-barrett May 19, 2026
e2528b6

Layout fix on split functions


gabriel-barrett May 19, 2026
3fbc821

Per group circuits


gabriel-barrett May 19, 2026
fefc240

`String` -> `Arc<str>`


gabriel-barrett May 19, 2026
8de0ae2

`address_eq` fast and slow paths


gabriel-barrett May 20, 2026
876e06c

Move split + layout from Rust to Lean


arthurpaulino May 20, 2026
b529b2a

Index return groups by USize at bytecode level


arthurpaulino May 20, 2026
42c4ceb

Use named structures for QueryCounts


arthurpaulino May 20, 2026
60c38d4

Preserve groupNames through dedup + match arms


arthurpaulino May 20, 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


1 change: 1 addition & 0 deletions
Ix/Aiur.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,5 +20,6 @@ public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Layout
public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Split
public import Ix.Aiur.Compiler
public import Ix.Aiur.Statistics

6 changes: 4 additions & 2 deletions
Ix/Aiur/Compiler.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@ public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Simple
public import Ix.Aiur.Compiler.Split

/-!
Aiur compiler pipeline: type-check, simplify, concretize, lower, deduplicate.
Expand DownExpand Up@@ -68,7 +69,7 @@ def Bytecode.Ctrl.collectConstrainedCallees (c : Bytecode.Ctrl) :
 | some block => branchCallees ++ block.collectConstrainedCallees
 | none => branchCallees
 withDefault ++ continuation.collectConstrainedCallees
 | .return _ _ | .yield _ _ => #[]
 | .return _ _ _ | .yield _ _ => #[]
termination_by (sizeOf c, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -117,9 +118,10 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev
 let (bytecodeRaw, preNameMap) ← concDecls.toBytecode
 let (bytecodeDedup, remap) := bytecodeRaw.deduplicate
 let needs := bytecodeDedup.needsCircuit
 let bytecode := { bytecodeDedup with
 let bytecodeConstrained : Bytecode.Toplevel := { bytecodeDedup with
 functions := bytecodeDedup.functions.mapIdx fun i f =>
 { f with constrained := needs[i]! } }
 let bytecode := bytecodeConstrained.computeFiltered
 let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx))
 fun acc name idx => acc.insert name (remap idx)
 pure (CompiledToplevel.mk t bytecode nameMap)
Expand Down

5 changes: 5 additions & 0 deletions
Ix/Aiur/Compiler/Check.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -790,6 +790,9 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
 | some sub => do pure (some (← inferTerm sub))
 let ret' ← inferTerm ret
 pure (Typed.Term.debug ret'.typ ret'.escapes label term' ret')
 | .retGroup name inner => do
 let inner' ← inferTerm inner
 pure (Typed.Term.retGroup inner'.typ inner'.escapes name inner')
termination_by (sizeOf t, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -917,6 +920,8 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
 | none => pure none
 | some sub => do pure (some (← zonkTypedTerm sub))
 pure (.debug (← zonkTyp τ) e label t' (← zonkTypedTerm r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← zonkTyp τ) e name (← zonkTypedTerm inner))
termination_by sizeOf t
decreasing_by
 all_goals first
Expand Down

8 changes: 8 additions & 0 deletions
Ix/Aiur/Compiler/Concretize.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -349,6 +349,8 @@ def termToConcrete
 | none => pure none
 | some sub => do pure (some (← termToConcrete mono sub))
 pure (.debug (← typToConcrete mono τ) e l t' (← termToConcrete mono r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← typToConcrete mono τ) e name (← termToConcrete mono inner))
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -541,6 +543,8 @@ def rewriteTypedTerm (decls : Typed.Decls)
 | none => none
 | some sub => some (rewriteTypedTerm decls subst mono sub)
 .debug (rewriteTyp subst mono τ) e l t' (rewriteTypedTerm decls subst mono r)
 | .retGroup τ e name inner =>
 .retGroup (rewriteTyp subst mono τ) e name (rewriteTypedTerm decls subst mono inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -625,6 +629,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
 let seen := collectInTyp seen τ
 let seen := match t with | some t => collectInTypedTerm seen t | none => seen
 collectInTypedTerm seen r
 | .retGroup τ _ _ inner => collectInTypedTerm (collectInTyp seen τ) inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -683,6 +688,7 @@ def collectCalls (decls : Typed.Decls)
 | .debug _ _ _ t r =>
 let seen := match t with | some t => collectCalls decls seen t | none => seen
 collectCalls decls seen r
 | .retGroup _ _ _ inner => collectCalls decls seen inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -771,6 +777,8 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
 | none => none
 | some sub => some (substInTypedTerm subst sub)
 .debug (Typ.instantiate subst τ) e l t' (substInTypedTerm subst r)
 | .retGroup τ e name inner =>
 .retGroup (Typ.instantiate subst τ) e name (substInTypedTerm subst inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand Down

7 changes: 4 additions & 3 deletions
Ix/Aiur/Compiler/Dedup.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -31,7 +31,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
 (match def_ with | none => none | some b => some (skeletonBlock b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
Expand DownExpand Up@@ -107,7 +107,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
 (match def_ with | none => none | some b => some (rewriteBlock f b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
Expand DownExpand Up@@ -196,7 +196,8 @@ def deduplicate_newFunctions (functions : Array Function) (classes : Array Nat)
 if can then
 let entry := deduplicate_class_entry functions classes cls
 let body := rewriteBlock remapFn f.body
 acc.push { body, layout := f.layout, entry, constrained := false }
 acc.push { body, layout := f.layout, groupNames := f.groupNames,
 entry, constrained := false }
 else acc)
 #[]

Expand Down

74 changes: 61 additions & 13 deletions
Ix/Aiur/Compiler/Lower.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -94,6 +94,15 @@ structure CompilerState where
 ops : Array Bytecode.Op
 selIdx : Bytecode.SelIdx
 degrees : Array Nat
 /-- Top of the `#[return_group(…)]` annotation stack — the display name
 whose index will tag every `Ctrl.return` emitted inside its scope. The
 index is looked up (or allocated) lazily at emit time. -/
 currentReturnGroupName : String := ""
 /-- Group display names allocated so far; position `i` is the name for
 group index `i`. -/
 groupNames : Array String := #[]
 /-- Inverse of `groupNames`: maps name → allocated index. -/
 groupNameMap : Std.HashMap String USize := {}
 deriving Inhabited

abbrev CompileM := EStateM String CompilerState
Expand DownExpand Up@@ -122,6 +131,20 @@ def pushOp (op : Bytecode.Op) (size : Nat := 1) : CompileM (Array Bytecode.ValId
def extractOps : CompileM (Array Bytecode.Op) :=
 modifyGet fun s => (s.ops, {s with ops := #[]})

/-- Look up the `USize` index for the current return-group name, allocating
fresh storage in `groupNames`/`groupNameMap` on first encounter. -/
def allocCurrentGroup : CompileM USize := do
 let st ← get
 let name := st.currentReturnGroupName
 match st.groupNameMap[name]? with
 | some idx => pure idx
 | none =>
 let idx : USize := USize.ofNat st.groupNames.size
 modify fun s => { s with
 groupNameMap := s.groupNameMap.insert name idx
 groupNames := s.groupNames.push name }
 pure idx

open Concrete in
mutual

Expand DownExpand Up@@ -299,6 +322,7 @@ def toIndex
 | some sub => do pure (some (← toIndex layoutMap bindings sub))
 modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
 toIndex layoutMap bindings ret
 | .retGroup _ _ _ inner => toIndex layoutMap bindings inner
termination_by (sizeOf term, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -446,6 +470,12 @@ def Concrete.Term.compile
 let data ← toIndex layoutMap bindings data
 modify fun stt => { stt with ops := stt.ops.push (.ioWrite data) }
 ret.compile returnTyp layoutMap bindings yieldCtrl
 | .retGroup _ _ name inner => do
 let oldGroup := (← get).currentReturnGroupName
 modify fun s => { s with currentReturnGroupName := name }
 let blk ← inner.compile returnTyp layoutMap bindings yieldCtrl
 modify fun s => { s with currentReturnGroupName := oldGroup }
 pure blk
 | .match _ _ scrut cases defaultOpt => do
 let idxs := bindings[scrut]?.getD #[0]
 let ops ← extractOps
Expand All@@ -460,21 +490,24 @@ def Concrete.Term.compile
 pure ({ ops, ctrl } : Bytecode.Block)
 | .ret _ _ term => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 pure ({ ops, ctrl := .return (id - 1) idxs } : Bytecode.Block)
 pure ({ ops, ctrl := .return (id - 1) groupIdx idxs } : Bytecode.Block)
 | _ => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 let ctrl : Bytecode.Ctrl :=
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs else .return (id - 1) idxs
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs
 else .return (id - 1) groupIdx idxs
 pure ({ ops, ctrl } : Bytecode.Block)
termination_by (sizeOf term, 0)
decreasing_by
Expand All@@ -497,7 +530,10 @@ def Concrete.addCase
 | .field g => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (g, term), defaultBlock)
 | .ref global pats => do
 let (index, offsets) ← match layoutMap[global]? with
Expand All@@ -516,22 +552,30 @@ def Concrete.addCase
 acc.insert patLocal slice
 let initState ← get
 let term ← term.compile returnTyp layoutMap ptrBindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (.ofNat index, term), defaultBlock)
 | .wildcard => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases, .some term)
 | _ => throw "addCase: unsupported pattern in concrete lower"
termination_by _ pair => (sizeOf pair.snd, 1)
decreasing_by all_goals first | decreasing_tactic | grind

end

/-- Lower a full concrete function to bytecode. -/
/-- Lower a full concrete function to bytecode. Returns the body, layout
state, and the per-function `groupNames` table (position `i` is the display
name for group index `i`). -/
def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 Except String (Bytecode.Block × Bytecode.LayoutMState) := do
 Except String (Bytecode.Block × Bytecode.LayoutMState × Array String) := do
 let (_inputSize, _outputSize) ← match layoutMap[f.name]? with
 | some (.function layout) => pure (layout.inputSize, layout.outputSize)
 | _ => throw s!"`{f.name}` should be a function"
Expand All@@ -542,15 +586,16 @@ def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 | .ok len => pure len
 let indices := Array.range' valIdx len
 pure (valIdx + len, bindings.insert arg indices)
 let state := { valIdx, selIdx := 0, ops := #[], degrees := Array.replicate valIdx 1 }
 let state : CompilerState := { valIdx, selIdx := 0, ops := #[],
 degrees := Array.replicate valIdx 1 }
 match f.body.compile f.output layoutMap bindings |>.run state with
 | .error e _ => throw e
 | .ok body _ =>
 | .ok body finalState =>
 let (_, layoutMState) := Bytecode.blockLayout body |>.run (.new valIdx)
 let layoutMState := { layoutMState with functionLayout :=
 { layoutMState.functionLayout with
 lookups := layoutMState.functionLayout.lookups + 1 } }
 pure (body, layoutMState)
 pure (body, layoutMState, finalState.groupNames)

def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 Except String (Bytecode.Toplevel × Std.HashMap Global Bytecode.FunIdx) := do
Expand All@@ -559,13 +604,16 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 let (functions, memSizes, nameMap) ← decls.foldlM (init := (#[], initMemSizes, {}))
 fun acc@(functions, memSizes, nameMap) (_, decl) => match decl with
 | .function function => do
 let (body, layoutMState) ← function.compile layout
 let (body, layoutMState, groupNames) ← function.compile layout
 let nameMap := nameMap.insert function.name functions.size
 let function := ⟨body, layoutMState.functionLayout, function.entry, false⟩
 let groupNames := if groupNames.isEmpty then #[""] else groupNames
 let function : Bytecode.Function :=
 { body, layout := layoutMState.functionLayout,
 groupNames, entry := function.entry, constrained := false }
 let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes
 pure (functions.push function, memSizes, nameMap)
 | _ => pure acc
 pure (⟨functions, memSizes.toArray⟩, nameMap)
 pure ({ functions, memorySizes := memSizes.toArray : Bytecode.Toplevel }, nameMap)

end Aiur

Expand Down

1 change: 1 addition & 0 deletions
Ix/Aiur/Compiler/Match.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -389,6 +389,7 @@ def typedToSimple : Term → Simple.Term
 | .debug τ e l t r =>
 let t' := match t with | none => none | some sub => some (typedToSimple sub)
 .debug τ e l t' (typedToSimple r)
 | .retGroup τ e name inner => .retGroup τ e name (typedToSimple inner)
termination_by t => sizeOf t
decreasing_by all_goals first | decreasing_tactic | grind

Expand Down

3 changes: 3 additions & 0 deletions
Ix/Aiur/Compiler/Simple.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -108,6 +108,9 @@ def simplifyTypedTerm (decls : Source.Decls) : Term → Except CheckError Term
 let a' ← simplifyTypedTerm decls a
 let b' ← simplifyTypedTerm decls b
 pure (.u32LessThan τ e a' b')
 | .retGroup τ e name inner => do
 let inner' ← simplifyTypedTerm decls inner
 pure (.retGroup τ e name inner')
 | t => pure t
termination_by t => sizeOf t
decreasing_by
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Force GitHub README to respect dark mode\n(function() {\n var style = document.createElement('style');\n style.textContent = '\n .markdown-body {\n color-scheme: dark light;\n }\n .markdown-body pre { background: #161b22 !important; }\n .markdown-body code { background: rgba(110, 118, 129, 0.4) !important; }\n .markdown-body table th, .markdown-body table td { border-color: #30363d !important; }\n .markdown-body img { background: #0d1117; }\n .markdown-body blockquote { border-left-color: #8b949e; }\n .markdown-body hr { border-color: #30363d; }\n ';\n document.head.appendChild(style);\n})();", "GitHub Dark Mode README Fix");
}
} catch(__e) { console.warn('[Userscript:GitHub Dark Mode README Fix]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Return groups#418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

gabriel-barrett
wants to merge
13
commits into

base:main
Choose a base branch

from

 Draft


Return groups
#418
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


13 commits


Select commit
Hold shift + click to select a range
be60c40

Return groups added to Source


gabriel-barrett May 19, 2026
70dcbfb

Return groups added to other states


gabriel-barrett May 19, 2026
843befa

Added groups to Lean `Bytecode.return`


gabriel-barrett May 19, 2026
a2a7b7f

Added return groups to `bytecode.rs`


gabriel-barrett May 19, 2026
5f7559a

Implemented `Function.split`


gabriel-barrett May 19, 2026
e2528b6

Layout fix on split functions


gabriel-barrett May 19, 2026
3fbc821

Per group circuits


gabriel-barrett May 19, 2026
fefc240

`String` -> `Arc<str>`


gabriel-barrett May 19, 2026
8de0ae2

`address_eq` fast and slow paths


gabriel-barrett May 20, 2026
876e06c

Move split + layout from Rust to Lean


arthurpaulino May 20, 2026
b529b2a

Index return groups by USize at bytecode level


arthurpaulino May 20, 2026
42c4ceb

Use named structures for QueryCounts


arthurpaulino May 20, 2026
60c38d4

Preserve groupNames through dedup + match arms


arthurpaulino May 20, 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


1 change: 1 addition & 0 deletions
Ix/Aiur.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,5 +20,6 @@ public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Layout
public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Split
public import Ix.Aiur.Compiler
public import Ix.Aiur.Statistics

6 changes: 4 additions & 2 deletions
Ix/Aiur/Compiler.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@ public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Simple
public import Ix.Aiur.Compiler.Split

/-!
Aiur compiler pipeline: type-check, simplify, concretize, lower, deduplicate.
Expand DownExpand Up@@ -68,7 +69,7 @@ def Bytecode.Ctrl.collectConstrainedCallees (c : Bytecode.Ctrl) :
 | some block => branchCallees ++ block.collectConstrainedCallees
 | none => branchCallees
 withDefault ++ continuation.collectConstrainedCallees
 | .return _ _ | .yield _ _ => #[]
 | .return _ _ _ | .yield _ _ => #[]
termination_by (sizeOf c, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -117,9 +118,10 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev
 let (bytecodeRaw, preNameMap) ← concDecls.toBytecode
 let (bytecodeDedup, remap) := bytecodeRaw.deduplicate
 let needs := bytecodeDedup.needsCircuit
 let bytecode := { bytecodeDedup with
 let bytecodeConstrained : Bytecode.Toplevel := { bytecodeDedup with
 functions := bytecodeDedup.functions.mapIdx fun i f =>
 { f with constrained := needs[i]! } }
 let bytecode := bytecodeConstrained.computeFiltered
 let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx))
 fun acc name idx => acc.insert name (remap idx)
 pure (CompiledToplevel.mk t bytecode nameMap)
Expand Down

5 changes: 5 additions & 0 deletions
Ix/Aiur/Compiler/Check.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -790,6 +790,9 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
 | some sub => do pure (some (← inferTerm sub))
 let ret' ← inferTerm ret
 pure (Typed.Term.debug ret'.typ ret'.escapes label term' ret')
 | .retGroup name inner => do
 let inner' ← inferTerm inner
 pure (Typed.Term.retGroup inner'.typ inner'.escapes name inner')
termination_by (sizeOf t, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -917,6 +920,8 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
 | none => pure none
 | some sub => do pure (some (← zonkTypedTerm sub))
 pure (.debug (← zonkTyp τ) e label t' (← zonkTypedTerm r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← zonkTyp τ) e name (← zonkTypedTerm inner))
termination_by sizeOf t
decreasing_by
 all_goals first
Expand Down

8 changes: 8 additions & 0 deletions
Ix/Aiur/Compiler/Concretize.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -349,6 +349,8 @@ def termToConcrete
 | none => pure none
 | some sub => do pure (some (← termToConcrete mono sub))
 pure (.debug (← typToConcrete mono τ) e l t' (← termToConcrete mono r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← typToConcrete mono τ) e name (← termToConcrete mono inner))
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -541,6 +543,8 @@ def rewriteTypedTerm (decls : Typed.Decls)
 | none => none
 | some sub => some (rewriteTypedTerm decls subst mono sub)
 .debug (rewriteTyp subst mono τ) e l t' (rewriteTypedTerm decls subst mono r)
 | .retGroup τ e name inner =>
 .retGroup (rewriteTyp subst mono τ) e name (rewriteTypedTerm decls subst mono inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -625,6 +629,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
 let seen := collectInTyp seen τ
 let seen := match t with | some t => collectInTypedTerm seen t | none => seen
 collectInTypedTerm seen r
 | .retGroup τ _ _ inner => collectInTypedTerm (collectInTyp seen τ) inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -683,6 +688,7 @@ def collectCalls (decls : Typed.Decls)
 | .debug _ _ _ t r =>
 let seen := match t with | some t => collectCalls decls seen t | none => seen
 collectCalls decls seen r
 | .retGroup _ _ _ inner => collectCalls decls seen inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -771,6 +777,8 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
 | none => none
 | some sub => some (substInTypedTerm subst sub)
 .debug (Typ.instantiate subst τ) e l t' (substInTypedTerm subst r)
 | .retGroup τ e name inner =>
 .retGroup (Typ.instantiate subst τ) e name (substInTypedTerm subst inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand Down

7 changes: 4 additions & 3 deletions
Ix/Aiur/Compiler/Dedup.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -31,7 +31,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
 (match def_ with | none => none | some b => some (skeletonBlock b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
Expand DownExpand Up@@ -107,7 +107,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
 (match def_ with | none => none | some b => some (rewriteBlock f b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
Expand DownExpand Up@@ -196,7 +196,8 @@ def deduplicate_newFunctions (functions : Array Function) (classes : Array Nat)
 if can then
 let entry := deduplicate_class_entry functions classes cls
 let body := rewriteBlock remapFn f.body
 acc.push { body, layout := f.layout, entry, constrained := false }
 acc.push { body, layout := f.layout, groupNames := f.groupNames,
 entry, constrained := false }
 else acc)
 #[]

Expand Down

74 changes: 61 additions & 13 deletions
Ix/Aiur/Compiler/Lower.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -94,6 +94,15 @@ structure CompilerState where
 ops : Array Bytecode.Op
 selIdx : Bytecode.SelIdx
 degrees : Array Nat
 /-- Top of the `#[return_group(…)]` annotation stack — the display name
 whose index will tag every `Ctrl.return` emitted inside its scope. The
 index is looked up (or allocated) lazily at emit time. -/
 currentReturnGroupName : String := ""
 /-- Group display names allocated so far; position `i` is the name for
 group index `i`. -/
 groupNames : Array String := #[]
 /-- Inverse of `groupNames`: maps name → allocated index. -/
 groupNameMap : Std.HashMap String USize := {}
 deriving Inhabited

abbrev CompileM := EStateM String CompilerState
Expand DownExpand Up@@ -122,6 +131,20 @@ def pushOp (op : Bytecode.Op) (size : Nat := 1) : CompileM (Array Bytecode.ValId
def extractOps : CompileM (Array Bytecode.Op) :=
 modifyGet fun s => (s.ops, {s with ops := #[]})

/-- Look up the `USize` index for the current return-group name, allocating
fresh storage in `groupNames`/`groupNameMap` on first encounter. -/
def allocCurrentGroup : CompileM USize := do
 let st ← get
 let name := st.currentReturnGroupName
 match st.groupNameMap[name]? with
 | some idx => pure idx
 | none =>
 let idx : USize := USize.ofNat st.groupNames.size
 modify fun s => { s with
 groupNameMap := s.groupNameMap.insert name idx
 groupNames := s.groupNames.push name }
 pure idx

open Concrete in
mutual

Expand DownExpand Up@@ -299,6 +322,7 @@ def toIndex
 | some sub => do pure (some (← toIndex layoutMap bindings sub))
 modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
 toIndex layoutMap bindings ret
 | .retGroup _ _ _ inner => toIndex layoutMap bindings inner
termination_by (sizeOf term, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -446,6 +470,12 @@ def Concrete.Term.compile
 let data ← toIndex layoutMap bindings data
 modify fun stt => { stt with ops := stt.ops.push (.ioWrite data) }
 ret.compile returnTyp layoutMap bindings yieldCtrl
 | .retGroup _ _ name inner => do
 let oldGroup := (← get).currentReturnGroupName
 modify fun s => { s with currentReturnGroupName := name }
 let blk ← inner.compile returnTyp layoutMap bindings yieldCtrl
 modify fun s => { s with currentReturnGroupName := oldGroup }
 pure blk
 | .match _ _ scrut cases defaultOpt => do
 let idxs := bindings[scrut]?.getD #[0]
 let ops ← extractOps
Expand All@@ -460,21 +490,24 @@ def Concrete.Term.compile
 pure ({ ops, ctrl } : Bytecode.Block)
 | .ret _ _ term => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 pure ({ ops, ctrl := .return (id - 1) idxs } : Bytecode.Block)
 pure ({ ops, ctrl := .return (id - 1) groupIdx idxs } : Bytecode.Block)
 | _ => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 let ctrl : Bytecode.Ctrl :=
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs else .return (id - 1) idxs
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs
 else .return (id - 1) groupIdx idxs
 pure ({ ops, ctrl } : Bytecode.Block)
termination_by (sizeOf term, 0)
decreasing_by
Expand All@@ -497,7 +530,10 @@ def Concrete.addCase
 | .field g => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (g, term), defaultBlock)
 | .ref global pats => do
 let (index, offsets) ← match layoutMap[global]? with
Expand All@@ -516,22 +552,30 @@ def Concrete.addCase
 acc.insert patLocal slice
 let initState ← get
 let term ← term.compile returnTyp layoutMap ptrBindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (.ofNat index, term), defaultBlock)
 | .wildcard => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases, .some term)
 | _ => throw "addCase: unsupported pattern in concrete lower"
termination_by _ pair => (sizeOf pair.snd, 1)
decreasing_by all_goals first | decreasing_tactic | grind

end

/-- Lower a full concrete function to bytecode. -/
/-- Lower a full concrete function to bytecode. Returns the body, layout
state, and the per-function `groupNames` table (position `i` is the display
name for group index `i`). -/
def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 Except String (Bytecode.Block × Bytecode.LayoutMState) := do
 Except String (Bytecode.Block × Bytecode.LayoutMState × Array String) := do
 let (_inputSize, _outputSize) ← match layoutMap[f.name]? with
 | some (.function layout) => pure (layout.inputSize, layout.outputSize)
 | _ => throw s!"`{f.name}` should be a function"
Expand All@@ -542,15 +586,16 @@ def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 | .ok len => pure len
 let indices := Array.range' valIdx len
 pure (valIdx + len, bindings.insert arg indices)
 let state := { valIdx, selIdx := 0, ops := #[], degrees := Array.replicate valIdx 1 }
 let state : CompilerState := { valIdx, selIdx := 0, ops := #[],
 degrees := Array.replicate valIdx 1 }
 match f.body.compile f.output layoutMap bindings |>.run state with
 | .error e _ => throw e
 | .ok body _ =>
 | .ok body finalState =>
 let (_, layoutMState) := Bytecode.blockLayout body |>.run (.new valIdx)
 let layoutMState := { layoutMState with functionLayout :=
 { layoutMState.functionLayout with
 lookups := layoutMState.functionLayout.lookups + 1 } }
 pure (body, layoutMState)
 pure (body, layoutMState, finalState.groupNames)

def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 Except String (Bytecode.Toplevel × Std.HashMap Global Bytecode.FunIdx) := do
Expand All@@ -559,13 +604,16 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 let (functions, memSizes, nameMap) ← decls.foldlM (init := (#[], initMemSizes, {}))
 fun acc@(functions, memSizes, nameMap) (_, decl) => match decl with
 | .function function => do
 let (body, layoutMState) ← function.compile layout
 let (body, layoutMState, groupNames) ← function.compile layout
 let nameMap := nameMap.insert function.name functions.size
 let function := ⟨body, layoutMState.functionLayout, function.entry, false⟩
 let groupNames := if groupNames.isEmpty then #[""] else groupNames
 let function : Bytecode.Function :=
 { body, layout := layoutMState.functionLayout,
 groupNames, entry := function.entry, constrained := false }
 let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes
 pure (functions.push function, memSizes, nameMap)
 | _ => pure acc
 pure (⟨functions, memSizes.toArray⟩, nameMap)
 pure ({ functions, memorySizes := memSizes.toArray : Bytecode.Toplevel }, nameMap)

end Aiur

Expand Down

1 change: 1 addition & 0 deletions
Ix/Aiur/Compiler/Match.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -389,6 +389,7 @@ def typedToSimple : Term → Simple.Term
 | .debug τ e l t r =>
 let t' := match t with | none => none | some sub => some (typedToSimple sub)
 .debug τ e l t' (typedToSimple r)
 | .retGroup τ e name inner => .retGroup τ e name (typedToSimple inner)
termination_by t => sizeOf t
decreasing_by all_goals first | decreasing_tactic | grind

Expand Down

3 changes: 3 additions & 0 deletions
Ix/Aiur/Compiler/Simple.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -108,6 +108,9 @@ def simplifyTypedTerm (decls : Source.Decls) : Term → Except CheckError Term
 let a' ← simplifyTypedTerm decls a
 let b' ← simplifyTypedTerm decls b
 pure (.u32LessThan τ e a' b')
 | .retGroup τ e name inner => do
 let inner' ← simplifyTypedTerm decls inner
 pure (.retGroup τ e name inner')
 | t => pure t
termination_by t => sizeOf t
decreasing_by
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Highlight search terms from Google/DuckDuckGo/Bing referrer\n(function() {\n var ref = document.referrer;\n var terms = [];\n \n if (ref.includes('google.com') || ref.includes('duckduckgo.com') || ref.includes('bing.com')) {\n var url = new URL(ref);\n var q = url.searchParams.get('q') || url.searchParams.get('p');\n if (q) {\n terms = q.split(/\\s+/).filter(function(t) { return t.length > 2; });\n }\n }\n \n if (terms.length === 0) return;\n \n var style = document.createElement('style');\n style.textContent = '.userscript-highlight { background: #fbbf24; color: #1a1a2e; padding: 1px 3px; border-radius: 2px; }';\n document.head.appendChild(style);\n \n function highlight(node) {\n if (node.nodeType === 3) { // text node\n var text = node.textContent;\n var found = false;\n terms.forEach(function(term) {\n var regex = new RegExp('(' + term.replace(/[.*+?^${}()|[\\]\\\\]/g, '\\\\') + ')', 'gi');\n if (regex.test(text)) {\n found = true;\n var frag = document.createDocumentFragment();\n var parts = text.split(regex);\n parts.forEach(function(part, i) {\n if (i % 2 === 0) {\n frag.appendChild(document.createTextNode(part));\n } else {\n var span = document.createElement('span');\n span.className = 'userscript-highlight';\n span.textContent = part;\n frag.appendChild(span);\n }\n });\n node.parentNode.replaceChild(frag, node);\n }\n });\n } else if (node.nodeType === 1 && node.childNodes) { // element\n var skipTags = ['SCRIPT', 'STYLE', 'NOSCRIPT', 'TEXTAREA', 'INPUT', 'SELECT'];\n if (!skipTags.includes(node.tagName)) {\n Array.from(node.childNodes).forEach(highlight);\n }\n }\n }\n \n highlight(document.body);\n \n // Re-highlight on dynamic content\n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1 || node.nodeType === 3) highlight(node);\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Highlight Search Terms");
}
} catch(__e) { console.warn('[Userscript:Highlight Search Terms]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Return groups#418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

gabriel-barrett
wants to merge
13
commits into

base:main
Choose a base branch

from

 Draft


Return groups
#418
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


13 commits


Select commit
Hold shift + click to select a range
be60c40

Return groups added to Source


gabriel-barrett May 19, 2026
70dcbfb

Return groups added to other states


gabriel-barrett May 19, 2026
843befa

Added groups to Lean `Bytecode.return`


gabriel-barrett May 19, 2026
a2a7b7f

Added return groups to `bytecode.rs`


gabriel-barrett May 19, 2026
5f7559a

Implemented `Function.split`


gabriel-barrett May 19, 2026
e2528b6

Layout fix on split functions


gabriel-barrett May 19, 2026
3fbc821

Per group circuits


gabriel-barrett May 19, 2026
fefc240

`String` -> `Arc<str>`


gabriel-barrett May 19, 2026
8de0ae2

`address_eq` fast and slow paths


gabriel-barrett May 20, 2026
876e06c

Move split + layout from Rust to Lean


arthurpaulino May 20, 2026
b529b2a

Index return groups by USize at bytecode level


arthurpaulino May 20, 2026
42c4ceb

Use named structures for QueryCounts


arthurpaulino May 20, 2026
60c38d4

Preserve groupNames through dedup + match arms


arthurpaulino May 20, 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


1 change: 1 addition & 0 deletions
Ix/Aiur.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,5 +20,6 @@ public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Layout
public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Split
public import Ix.Aiur.Compiler
public import Ix.Aiur.Statistics

6 changes: 4 additions & 2 deletions
Ix/Aiur/Compiler.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@ public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Simple
public import Ix.Aiur.Compiler.Split

/-!
Aiur compiler pipeline: type-check, simplify, concretize, lower, deduplicate.
Expand DownExpand Up@@ -68,7 +69,7 @@ def Bytecode.Ctrl.collectConstrainedCallees (c : Bytecode.Ctrl) :
 | some block => branchCallees ++ block.collectConstrainedCallees
 | none => branchCallees
 withDefault ++ continuation.collectConstrainedCallees
 | .return _ _ | .yield _ _ => #[]
 | .return _ _ _ | .yield _ _ => #[]
termination_by (sizeOf c, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -117,9 +118,10 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev
 let (bytecodeRaw, preNameMap) ← concDecls.toBytecode
 let (bytecodeDedup, remap) := bytecodeRaw.deduplicate
 let needs := bytecodeDedup.needsCircuit
 let bytecode := { bytecodeDedup with
 let bytecodeConstrained : Bytecode.Toplevel := { bytecodeDedup with
 functions := bytecodeDedup.functions.mapIdx fun i f =>
 { f with constrained := needs[i]! } }
 let bytecode := bytecodeConstrained.computeFiltered
 let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx))
 fun acc name idx => acc.insert name (remap idx)
 pure (CompiledToplevel.mk t bytecode nameMap)
Expand Down

5 changes: 5 additions & 0 deletions
Ix/Aiur/Compiler/Check.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -790,6 +790,9 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
 | some sub => do pure (some (← inferTerm sub))
 let ret' ← inferTerm ret
 pure (Typed.Term.debug ret'.typ ret'.escapes label term' ret')
 | .retGroup name inner => do
 let inner' ← inferTerm inner
 pure (Typed.Term.retGroup inner'.typ inner'.escapes name inner')
termination_by (sizeOf t, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -917,6 +920,8 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
 | none => pure none
 | some sub => do pure (some (← zonkTypedTerm sub))
 pure (.debug (← zonkTyp τ) e label t' (← zonkTypedTerm r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← zonkTyp τ) e name (← zonkTypedTerm inner))
termination_by sizeOf t
decreasing_by
 all_goals first
Expand Down

8 changes: 8 additions & 0 deletions
Ix/Aiur/Compiler/Concretize.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -349,6 +349,8 @@ def termToConcrete
 | none => pure none
 | some sub => do pure (some (← termToConcrete mono sub))
 pure (.debug (← typToConcrete mono τ) e l t' (← termToConcrete mono r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← typToConcrete mono τ) e name (← termToConcrete mono inner))
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -541,6 +543,8 @@ def rewriteTypedTerm (decls : Typed.Decls)
 | none => none
 | some sub => some (rewriteTypedTerm decls subst mono sub)
 .debug (rewriteTyp subst mono τ) e l t' (rewriteTypedTerm decls subst mono r)
 | .retGroup τ e name inner =>
 .retGroup (rewriteTyp subst mono τ) e name (rewriteTypedTerm decls subst mono inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -625,6 +629,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
 let seen := collectInTyp seen τ
 let seen := match t with | some t => collectInTypedTerm seen t | none => seen
 collectInTypedTerm seen r
 | .retGroup τ _ _ inner => collectInTypedTerm (collectInTyp seen τ) inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -683,6 +688,7 @@ def collectCalls (decls : Typed.Decls)
 | .debug _ _ _ t r =>
 let seen := match t with | some t => collectCalls decls seen t | none => seen
 collectCalls decls seen r
 | .retGroup _ _ _ inner => collectCalls decls seen inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -771,6 +777,8 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
 | none => none
 | some sub => some (substInTypedTerm subst sub)
 .debug (Typ.instantiate subst τ) e l t' (substInTypedTerm subst r)
 | .retGroup τ e name inner =>
 .retGroup (Typ.instantiate subst τ) e name (substInTypedTerm subst inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand Down

7 changes: 4 additions & 3 deletions
Ix/Aiur/Compiler/Dedup.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -31,7 +31,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
 (match def_ with | none => none | some b => some (skeletonBlock b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
Expand DownExpand Up@@ -107,7 +107,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
 (match def_ with | none => none | some b => some (rewriteBlock f b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
Expand DownExpand Up@@ -196,7 +196,8 @@ def deduplicate_newFunctions (functions : Array Function) (classes : Array Nat)
 if can then
 let entry := deduplicate_class_entry functions classes cls
 let body := rewriteBlock remapFn f.body
 acc.push { body, layout := f.layout, entry, constrained := false }
 acc.push { body, layout := f.layout, groupNames := f.groupNames,
 entry, constrained := false }
 else acc)
 #[]

Expand Down

74 changes: 61 additions & 13 deletions
Ix/Aiur/Compiler/Lower.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -94,6 +94,15 @@ structure CompilerState where
 ops : Array Bytecode.Op
 selIdx : Bytecode.SelIdx
 degrees : Array Nat
 /-- Top of the `#[return_group(…)]` annotation stack — the display name
 whose index will tag every `Ctrl.return` emitted inside its scope. The
 index is looked up (or allocated) lazily at emit time. -/
 currentReturnGroupName : String := ""
 /-- Group display names allocated so far; position `i` is the name for
 group index `i`. -/
 groupNames : Array String := #[]
 /-- Inverse of `groupNames`: maps name → allocated index. -/
 groupNameMap : Std.HashMap String USize := {}
 deriving Inhabited

abbrev CompileM := EStateM String CompilerState
Expand DownExpand Up@@ -122,6 +131,20 @@ def pushOp (op : Bytecode.Op) (size : Nat := 1) : CompileM (Array Bytecode.ValId
def extractOps : CompileM (Array Bytecode.Op) :=
 modifyGet fun s => (s.ops, {s with ops := #[]})

/-- Look up the `USize` index for the current return-group name, allocating
fresh storage in `groupNames`/`groupNameMap` on first encounter. -/
def allocCurrentGroup : CompileM USize := do
 let st ← get
 let name := st.currentReturnGroupName
 match st.groupNameMap[name]? with
 | some idx => pure idx
 | none =>
 let idx : USize := USize.ofNat st.groupNames.size
 modify fun s => { s with
 groupNameMap := s.groupNameMap.insert name idx
 groupNames := s.groupNames.push name }
 pure idx

open Concrete in
mutual

Expand DownExpand Up@@ -299,6 +322,7 @@ def toIndex
 | some sub => do pure (some (← toIndex layoutMap bindings sub))
 modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
 toIndex layoutMap bindings ret
 | .retGroup _ _ _ inner => toIndex layoutMap bindings inner
termination_by (sizeOf term, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -446,6 +470,12 @@ def Concrete.Term.compile
 let data ← toIndex layoutMap bindings data
 modify fun stt => { stt with ops := stt.ops.push (.ioWrite data) }
 ret.compile returnTyp layoutMap bindings yieldCtrl
 | .retGroup _ _ name inner => do
 let oldGroup := (← get).currentReturnGroupName
 modify fun s => { s with currentReturnGroupName := name }
 let blk ← inner.compile returnTyp layoutMap bindings yieldCtrl
 modify fun s => { s with currentReturnGroupName := oldGroup }
 pure blk
 | .match _ _ scrut cases defaultOpt => do
 let idxs := bindings[scrut]?.getD #[0]
 let ops ← extractOps
Expand All@@ -460,21 +490,24 @@ def Concrete.Term.compile
 pure ({ ops, ctrl } : Bytecode.Block)
 | .ret _ _ term => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 pure ({ ops, ctrl := .return (id - 1) idxs } : Bytecode.Block)
 pure ({ ops, ctrl := .return (id - 1) groupIdx idxs } : Bytecode.Block)
 | _ => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 let ctrl : Bytecode.Ctrl :=
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs else .return (id - 1) idxs
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs
 else .return (id - 1) groupIdx idxs
 pure ({ ops, ctrl } : Bytecode.Block)
termination_by (sizeOf term, 0)
decreasing_by
Expand All@@ -497,7 +530,10 @@ def Concrete.addCase
 | .field g => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (g, term), defaultBlock)
 | .ref global pats => do
 let (index, offsets) ← match layoutMap[global]? with
Expand All@@ -516,22 +552,30 @@ def Concrete.addCase
 acc.insert patLocal slice
 let initState ← get
 let term ← term.compile returnTyp layoutMap ptrBindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (.ofNat index, term), defaultBlock)
 | .wildcard => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases, .some term)
 | _ => throw "addCase: unsupported pattern in concrete lower"
termination_by _ pair => (sizeOf pair.snd, 1)
decreasing_by all_goals first | decreasing_tactic | grind

end

/-- Lower a full concrete function to bytecode. -/
/-- Lower a full concrete function to bytecode. Returns the body, layout
state, and the per-function `groupNames` table (position `i` is the display
name for group index `i`). -/
def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 Except String (Bytecode.Block × Bytecode.LayoutMState) := do
 Except String (Bytecode.Block × Bytecode.LayoutMState × Array String) := do
 let (_inputSize, _outputSize) ← match layoutMap[f.name]? with
 | some (.function layout) => pure (layout.inputSize, layout.outputSize)
 | _ => throw s!"`{f.name}` should be a function"
Expand All@@ -542,15 +586,16 @@ def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 | .ok len => pure len
 let indices := Array.range' valIdx len
 pure (valIdx + len, bindings.insert arg indices)
 let state := { valIdx, selIdx := 0, ops := #[], degrees := Array.replicate valIdx 1 }
 let state : CompilerState := { valIdx, selIdx := 0, ops := #[],
 degrees := Array.replicate valIdx 1 }
 match f.body.compile f.output layoutMap bindings |>.run state with
 | .error e _ => throw e
 | .ok body _ =>
 | .ok body finalState =>
 let (_, layoutMState) := Bytecode.blockLayout body |>.run (.new valIdx)
 let layoutMState := { layoutMState with functionLayout :=
 { layoutMState.functionLayout with
 lookups := layoutMState.functionLayout.lookups + 1 } }
 pure (body, layoutMState)
 pure (body, layoutMState, finalState.groupNames)

def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 Except String (Bytecode.Toplevel × Std.HashMap Global Bytecode.FunIdx) := do
Expand All@@ -559,13 +604,16 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 let (functions, memSizes, nameMap) ← decls.foldlM (init := (#[], initMemSizes, {}))
 fun acc@(functions, memSizes, nameMap) (_, decl) => match decl with
 | .function function => do
 let (body, layoutMState) ← function.compile layout
 let (body, layoutMState, groupNames) ← function.compile layout
 let nameMap := nameMap.insert function.name functions.size
 let function := ⟨body, layoutMState.functionLayout, function.entry, false⟩
 let groupNames := if groupNames.isEmpty then #[""] else groupNames
 let function : Bytecode.Function :=
 { body, layout := layoutMState.functionLayout,
 groupNames, entry := function.entry, constrained := false }
 let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes
 pure (functions.push function, memSizes, nameMap)
 | _ => pure acc
 pure (⟨functions, memSizes.toArray⟩, nameMap)
 pure ({ functions, memorySizes := memSizes.toArray : Bytecode.Toplevel }, nameMap)

end Aiur

Expand Down

1 change: 1 addition & 0 deletions
Ix/Aiur/Compiler/Match.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -389,6 +389,7 @@ def typedToSimple : Term → Simple.Term
 | .debug τ e l t r =>
 let t' := match t with | none => none | some sub => some (typedToSimple sub)
 .debug τ e l t' (typedToSimple r)
 | .retGroup τ e name inner => .retGroup τ e name (typedToSimple inner)
termination_by t => sizeOf t
decreasing_by all_goals first | decreasing_tactic | grind

Expand Down

3 changes: 3 additions & 0 deletions
Ix/Aiur/Compiler/Simple.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -108,6 +108,9 @@ def simplifyTypedTerm (decls : Source.Decls) : Term → Except CheckError Term
 let a' ← simplifyTypedTerm decls a
 let b' ← simplifyTypedTerm decls b
 pure (.u32LessThan τ e a' b')
 | .retGroup τ e name inner => do
 let inner' ← simplifyTypedTerm decls inner
 pure (.retGroup τ e name inner')
 | t => pure t
termination_by t => sizeOf t
decreasing_by
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Strip utm_, fbclid, gclid, etc. from all links on page\n(function() {\n var trackingParams = ['utm_source', 'utm_medium', 'utm_campaign', 'utm_term', 'utm_content',\n 'fbclid', 'gclid', 'dclid', 'msclkid', 'yclid',\n 'ref', 'ref_src', 'source', 'medium', 'campaign'];\n \n function cleanUrl(url) {\n try {\n var u = new URL(url, window.location.origin);\n var changed = false;\n trackingParams.forEach(function(p) {\n if (u.searchParams.has(p)) {\n u.searchParams.delete(p);\n changed = true;\n }\n });\n return changed ? u.toString() : url;\n } catch (e) {\n return url;\n }\n }\n \n function cleanLinks() {\n document.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n \n cleanLinks();\n \n var observer = new MutationObserver(function(mutations) {\n mutations.forEach(function(m) {\n m.addedNodes.forEach(function(node) {\n if (node.nodeType === 1) {\n if (node.tagName === 'A') cleanLinks();\n node.querySelectorAll('a[href]').forEach(function(a) {\n var clean = cleanUrl(a.href);\n if (clean !== a.href) a.href = clean;\n });\n }\n });\n });\n });\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "Remove Tracking Parameters from Links");
}
} catch(__e) { console.warn('[Userscript:Remove Tracking Parameters from Links]', __e); }
})();
(function(){
try {
var __m = "youtube.com";
var __re = new RegExp('^' + "youtube\\.com" + '

Skip to content
Return groups#418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

gabriel-barrett
wants to merge
13
commits into

base:main
Choose a base branch

from

 Draft


Return groups
#418
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


13 commits


Select commit
Hold shift + click to select a range
be60c40

Return groups added to Source


gabriel-barrett May 19, 2026
70dcbfb

Return groups added to other states


gabriel-barrett May 19, 2026
843befa

Added groups to Lean `Bytecode.return`


gabriel-barrett May 19, 2026
a2a7b7f

Added return groups to `bytecode.rs`


gabriel-barrett May 19, 2026
5f7559a

Implemented `Function.split`


gabriel-barrett May 19, 2026
e2528b6

Layout fix on split functions


gabriel-barrett May 19, 2026
3fbc821

Per group circuits


gabriel-barrett May 19, 2026
fefc240

`String` -> `Arc<str>`


gabriel-barrett May 19, 2026
8de0ae2

`address_eq` fast and slow paths


gabriel-barrett May 20, 2026
876e06c

Move split + layout from Rust to Lean


arthurpaulino May 20, 2026
b529b2a

Index return groups by USize at bytecode level


arthurpaulino May 20, 2026
42c4ceb

Use named structures for QueryCounts


arthurpaulino May 20, 2026
60c38d4

Preserve groupNames through dedup + match arms


arthurpaulino May 20, 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


1 change: 1 addition & 0 deletions
Ix/Aiur.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,5 +20,6 @@ public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Layout
public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Split
public import Ix.Aiur.Compiler
public import Ix.Aiur.Statistics

6 changes: 4 additions & 2 deletions
Ix/Aiur/Compiler.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@ public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Simple
public import Ix.Aiur.Compiler.Split

/-!
Aiur compiler pipeline: type-check, simplify, concretize, lower, deduplicate.
Expand DownExpand Up@@ -68,7 +69,7 @@ def Bytecode.Ctrl.collectConstrainedCallees (c : Bytecode.Ctrl) :
 | some block => branchCallees ++ block.collectConstrainedCallees
 | none => branchCallees
 withDefault ++ continuation.collectConstrainedCallees
 | .return _ _ | .yield _ _ => #[]
 | .return _ _ _ | .yield _ _ => #[]
termination_by (sizeOf c, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -117,9 +118,10 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev
 let (bytecodeRaw, preNameMap) ← concDecls.toBytecode
 let (bytecodeDedup, remap) := bytecodeRaw.deduplicate
 let needs := bytecodeDedup.needsCircuit
 let bytecode := { bytecodeDedup with
 let bytecodeConstrained : Bytecode.Toplevel := { bytecodeDedup with
 functions := bytecodeDedup.functions.mapIdx fun i f =>
 { f with constrained := needs[i]! } }
 let bytecode := bytecodeConstrained.computeFiltered
 let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx))
 fun acc name idx => acc.insert name (remap idx)
 pure (CompiledToplevel.mk t bytecode nameMap)
Expand Down

5 changes: 5 additions & 0 deletions
Ix/Aiur/Compiler/Check.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -790,6 +790,9 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
 | some sub => do pure (some (← inferTerm sub))
 let ret' ← inferTerm ret
 pure (Typed.Term.debug ret'.typ ret'.escapes label term' ret')
 | .retGroup name inner => do
 let inner' ← inferTerm inner
 pure (Typed.Term.retGroup inner'.typ inner'.escapes name inner')
termination_by (sizeOf t, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -917,6 +920,8 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
 | none => pure none
 | some sub => do pure (some (← zonkTypedTerm sub))
 pure (.debug (← zonkTyp τ) e label t' (← zonkTypedTerm r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← zonkTyp τ) e name (← zonkTypedTerm inner))
termination_by sizeOf t
decreasing_by
 all_goals first
Expand Down

8 changes: 8 additions & 0 deletions
Ix/Aiur/Compiler/Concretize.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -349,6 +349,8 @@ def termToConcrete
 | none => pure none
 | some sub => do pure (some (← termToConcrete mono sub))
 pure (.debug (← typToConcrete mono τ) e l t' (← termToConcrete mono r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← typToConcrete mono τ) e name (← termToConcrete mono inner))
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -541,6 +543,8 @@ def rewriteTypedTerm (decls : Typed.Decls)
 | none => none
 | some sub => some (rewriteTypedTerm decls subst mono sub)
 .debug (rewriteTyp subst mono τ) e l t' (rewriteTypedTerm decls subst mono r)
 | .retGroup τ e name inner =>
 .retGroup (rewriteTyp subst mono τ) e name (rewriteTypedTerm decls subst mono inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -625,6 +629,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
 let seen := collectInTyp seen τ
 let seen := match t with | some t => collectInTypedTerm seen t | none => seen
 collectInTypedTerm seen r
 | .retGroup τ _ _ inner => collectInTypedTerm (collectInTyp seen τ) inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -683,6 +688,7 @@ def collectCalls (decls : Typed.Decls)
 | .debug _ _ _ t r =>
 let seen := match t with | some t => collectCalls decls seen t | none => seen
 collectCalls decls seen r
 | .retGroup _ _ _ inner => collectCalls decls seen inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -771,6 +777,8 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
 | none => none
 | some sub => some (substInTypedTerm subst sub)
 .debug (Typ.instantiate subst τ) e l t' (substInTypedTerm subst r)
 | .retGroup τ e name inner =>
 .retGroup (Typ.instantiate subst τ) e name (substInTypedTerm subst inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand Down

7 changes: 4 additions & 3 deletions
Ix/Aiur/Compiler/Dedup.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -31,7 +31,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
 (match def_ with | none => none | some b => some (skeletonBlock b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
Expand DownExpand Up@@ -107,7 +107,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
 (match def_ with | none => none | some b => some (rewriteBlock f b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
Expand DownExpand Up@@ -196,7 +196,8 @@ def deduplicate_newFunctions (functions : Array Function) (classes : Array Nat)
 if can then
 let entry := deduplicate_class_entry functions classes cls
 let body := rewriteBlock remapFn f.body
 acc.push { body, layout := f.layout, entry, constrained := false }
 acc.push { body, layout := f.layout, groupNames := f.groupNames,
 entry, constrained := false }
 else acc)
 #[]

Expand Down

74 changes: 61 additions & 13 deletions
Ix/Aiur/Compiler/Lower.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -94,6 +94,15 @@ structure CompilerState where
 ops : Array Bytecode.Op
 selIdx : Bytecode.SelIdx
 degrees : Array Nat
 /-- Top of the `#[return_group(…)]` annotation stack — the display name
 whose index will tag every `Ctrl.return` emitted inside its scope. The
 index is looked up (or allocated) lazily at emit time. -/
 currentReturnGroupName : String := ""
 /-- Group display names allocated so far; position `i` is the name for
 group index `i`. -/
 groupNames : Array String := #[]
 /-- Inverse of `groupNames`: maps name → allocated index. -/
 groupNameMap : Std.HashMap String USize := {}
 deriving Inhabited

abbrev CompileM := EStateM String CompilerState
Expand DownExpand Up@@ -122,6 +131,20 @@ def pushOp (op : Bytecode.Op) (size : Nat := 1) : CompileM (Array Bytecode.ValId
def extractOps : CompileM (Array Bytecode.Op) :=
 modifyGet fun s => (s.ops, {s with ops := #[]})

/-- Look up the `USize` index for the current return-group name, allocating
fresh storage in `groupNames`/`groupNameMap` on first encounter. -/
def allocCurrentGroup : CompileM USize := do
 let st ← get
 let name := st.currentReturnGroupName
 match st.groupNameMap[name]? with
 | some idx => pure idx
 | none =>
 let idx : USize := USize.ofNat st.groupNames.size
 modify fun s => { s with
 groupNameMap := s.groupNameMap.insert name idx
 groupNames := s.groupNames.push name }
 pure idx

open Concrete in
mutual

Expand DownExpand Up@@ -299,6 +322,7 @@ def toIndex
 | some sub => do pure (some (← toIndex layoutMap bindings sub))
 modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
 toIndex layoutMap bindings ret
 | .retGroup _ _ _ inner => toIndex layoutMap bindings inner
termination_by (sizeOf term, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -446,6 +470,12 @@ def Concrete.Term.compile
 let data ← toIndex layoutMap bindings data
 modify fun stt => { stt with ops := stt.ops.push (.ioWrite data) }
 ret.compile returnTyp layoutMap bindings yieldCtrl
 | .retGroup _ _ name inner => do
 let oldGroup := (← get).currentReturnGroupName
 modify fun s => { s with currentReturnGroupName := name }
 let blk ← inner.compile returnTyp layoutMap bindings yieldCtrl
 modify fun s => { s with currentReturnGroupName := oldGroup }
 pure blk
 | .match _ _ scrut cases defaultOpt => do
 let idxs := bindings[scrut]?.getD #[0]
 let ops ← extractOps
Expand All@@ -460,21 +490,24 @@ def Concrete.Term.compile
 pure ({ ops, ctrl } : Bytecode.Block)
 | .ret _ _ term => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 pure ({ ops, ctrl := .return (id - 1) idxs } : Bytecode.Block)
 pure ({ ops, ctrl := .return (id - 1) groupIdx idxs } : Bytecode.Block)
 | _ => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 let ctrl : Bytecode.Ctrl :=
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs else .return (id - 1) idxs
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs
 else .return (id - 1) groupIdx idxs
 pure ({ ops, ctrl } : Bytecode.Block)
termination_by (sizeOf term, 0)
decreasing_by
Expand All@@ -497,7 +530,10 @@ def Concrete.addCase
 | .field g => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (g, term), defaultBlock)
 | .ref global pats => do
 let (index, offsets) ← match layoutMap[global]? with
Expand All@@ -516,22 +552,30 @@ def Concrete.addCase
 acc.insert patLocal slice
 let initState ← get
 let term ← term.compile returnTyp layoutMap ptrBindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (.ofNat index, term), defaultBlock)
 | .wildcard => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases, .some term)
 | _ => throw "addCase: unsupported pattern in concrete lower"
termination_by _ pair => (sizeOf pair.snd, 1)
decreasing_by all_goals first | decreasing_tactic | grind

end

/-- Lower a full concrete function to bytecode. -/
/-- Lower a full concrete function to bytecode. Returns the body, layout
state, and the per-function `groupNames` table (position `i` is the display
name for group index `i`). -/
def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 Except String (Bytecode.Block × Bytecode.LayoutMState) := do
 Except String (Bytecode.Block × Bytecode.LayoutMState × Array String) := do
 let (_inputSize, _outputSize) ← match layoutMap[f.name]? with
 | some (.function layout) => pure (layout.inputSize, layout.outputSize)
 | _ => throw s!"`{f.name}` should be a function"
Expand All@@ -542,15 +586,16 @@ def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 | .ok len => pure len
 let indices := Array.range' valIdx len
 pure (valIdx + len, bindings.insert arg indices)
 let state := { valIdx, selIdx := 0, ops := #[], degrees := Array.replicate valIdx 1 }
 let state : CompilerState := { valIdx, selIdx := 0, ops := #[],
 degrees := Array.replicate valIdx 1 }
 match f.body.compile f.output layoutMap bindings |>.run state with
 | .error e _ => throw e
 | .ok body _ =>
 | .ok body finalState =>
 let (_, layoutMState) := Bytecode.blockLayout body |>.run (.new valIdx)
 let layoutMState := { layoutMState with functionLayout :=
 { layoutMState.functionLayout with
 lookups := layoutMState.functionLayout.lookups + 1 } }
 pure (body, layoutMState)
 pure (body, layoutMState, finalState.groupNames)

def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 Except String (Bytecode.Toplevel × Std.HashMap Global Bytecode.FunIdx) := do
Expand All@@ -559,13 +604,16 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 let (functions, memSizes, nameMap) ← decls.foldlM (init := (#[], initMemSizes, {}))
 fun acc@(functions, memSizes, nameMap) (_, decl) => match decl with
 | .function function => do
 let (body, layoutMState) ← function.compile layout
 let (body, layoutMState, groupNames) ← function.compile layout
 let nameMap := nameMap.insert function.name functions.size
 let function := ⟨body, layoutMState.functionLayout, function.entry, false⟩
 let groupNames := if groupNames.isEmpty then #[""] else groupNames
 let function : Bytecode.Function :=
 { body, layout := layoutMState.functionLayout,
 groupNames, entry := function.entry, constrained := false }
 let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes
 pure (functions.push function, memSizes, nameMap)
 | _ => pure acc
 pure (⟨functions, memSizes.toArray⟩, nameMap)
 pure ({ functions, memorySizes := memSizes.toArray : Bytecode.Toplevel }, nameMap)

end Aiur

Expand Down

1 change: 1 addition & 0 deletions
Ix/Aiur/Compiler/Match.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -389,6 +389,7 @@ def typedToSimple : Term → Simple.Term
 | .debug τ e l t r =>
 let t' := match t with | none => none | some sub => some (typedToSimple sub)
 .debug τ e l t' (typedToSimple r)
 | .retGroup τ e name inner => .retGroup τ e name (typedToSimple inner)
termination_by t => sizeOf t
decreasing_by all_goals first | decreasing_tactic | grind

Expand Down

3 changes: 3 additions & 0 deletions
Ix/Aiur/Compiler/Simple.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -108,6 +108,9 @@ def simplifyTypedTerm (decls : Source.Decls) : Term → Except CheckError Term
 let a' ← simplifyTypedTerm decls a
 let b' ← simplifyTypedTerm decls b
 pure (.u32LessThan τ e a' b')
 | .retGroup τ e name inner => do
 let inner' ← simplifyTypedTerm decls inner
 pure (.retGroup τ e name inner')
 | t => pure t
termination_by t => sizeOf t
decreasing_by
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Auto-enable theater mode on YouTube\n(function() {\n function tryTheater() {\n var btn = document.querySelector('button[aria-label=\"Theater mode\"], ytd-player #player button[title=\"Theater mode\"]');\n if (btn && !btn.classList.contains('activated')) {\n btn.click();\n }\n }\n \n // Try immediately\n tryTheater();\n \n // Try after navigation (SPA)\n var lastUrl = location.href;\n setInterval(function() {\n if (location.href !== lastUrl) {\n lastUrl = location.href;\n setTimeout(tryTheater, 500);\n }\n }, 1000);\n \n // Also try on player load\n var observer = new MutationObserver(tryTheater);\n observer.observe(document.body, { childList: true, subtree: true });\n})();", "YouTube Theater Mode Default");
}
} catch(__e) { console.warn('[Userscript:YouTube Theater Mode Default]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Return groups#418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

gabriel-barrett
wants to merge
13
commits into

base:main
Choose a base branch

from

 Draft


Return groups
#418
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


13 commits


Select commit
Hold shift + click to select a range
be60c40

Return groups added to Source


gabriel-barrett May 19, 2026
70dcbfb

Return groups added to other states


gabriel-barrett May 19, 2026
843befa

Added groups to Lean `Bytecode.return`


gabriel-barrett May 19, 2026
a2a7b7f

Added return groups to `bytecode.rs`


gabriel-barrett May 19, 2026
5f7559a

Implemented `Function.split`


gabriel-barrett May 19, 2026
e2528b6

Layout fix on split functions


gabriel-barrett May 19, 2026
3fbc821

Per group circuits


gabriel-barrett May 19, 2026
fefc240

`String` -> `Arc<str>`


gabriel-barrett May 19, 2026
8de0ae2

`address_eq` fast and slow paths


gabriel-barrett May 20, 2026
876e06c

Move split + layout from Rust to Lean


arthurpaulino May 20, 2026
b529b2a

Index return groups by USize at bytecode level


arthurpaulino May 20, 2026
42c4ceb

Use named structures for QueryCounts


arthurpaulino May 20, 2026
60c38d4

Preserve groupNames through dedup + match arms


arthurpaulino May 20, 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


1 change: 1 addition & 0 deletions
Ix/Aiur.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,5 +20,6 @@ public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Layout
public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Split
public import Ix.Aiur.Compiler
public import Ix.Aiur.Statistics

6 changes: 4 additions & 2 deletions
Ix/Aiur/Compiler.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@ public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Simple
public import Ix.Aiur.Compiler.Split

/-!
Aiur compiler pipeline: type-check, simplify, concretize, lower, deduplicate.
Expand DownExpand Up@@ -68,7 +69,7 @@ def Bytecode.Ctrl.collectConstrainedCallees (c : Bytecode.Ctrl) :
 | some block => branchCallees ++ block.collectConstrainedCallees
 | none => branchCallees
 withDefault ++ continuation.collectConstrainedCallees
 | .return _ _ | .yield _ _ => #[]
 | .return _ _ _ | .yield _ _ => #[]
termination_by (sizeOf c, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -117,9 +118,10 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev
 let (bytecodeRaw, preNameMap) ← concDecls.toBytecode
 let (bytecodeDedup, remap) := bytecodeRaw.deduplicate
 let needs := bytecodeDedup.needsCircuit
 let bytecode := { bytecodeDedup with
 let bytecodeConstrained : Bytecode.Toplevel := { bytecodeDedup with
 functions := bytecodeDedup.functions.mapIdx fun i f =>
 { f with constrained := needs[i]! } }
 let bytecode := bytecodeConstrained.computeFiltered
 let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx))
 fun acc name idx => acc.insert name (remap idx)
 pure (CompiledToplevel.mk t bytecode nameMap)
Expand Down

5 changes: 5 additions & 0 deletions
Ix/Aiur/Compiler/Check.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -790,6 +790,9 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
 | some sub => do pure (some (← inferTerm sub))
 let ret' ← inferTerm ret
 pure (Typed.Term.debug ret'.typ ret'.escapes label term' ret')
 | .retGroup name inner => do
 let inner' ← inferTerm inner
 pure (Typed.Term.retGroup inner'.typ inner'.escapes name inner')
termination_by (sizeOf t, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -917,6 +920,8 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
 | none => pure none
 | some sub => do pure (some (← zonkTypedTerm sub))
 pure (.debug (← zonkTyp τ) e label t' (← zonkTypedTerm r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← zonkTyp τ) e name (← zonkTypedTerm inner))
termination_by sizeOf t
decreasing_by
 all_goals first
Expand Down

8 changes: 8 additions & 0 deletions
Ix/Aiur/Compiler/Concretize.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -349,6 +349,8 @@ def termToConcrete
 | none => pure none
 | some sub => do pure (some (← termToConcrete mono sub))
 pure (.debug (← typToConcrete mono τ) e l t' (← termToConcrete mono r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← typToConcrete mono τ) e name (← termToConcrete mono inner))
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -541,6 +543,8 @@ def rewriteTypedTerm (decls : Typed.Decls)
 | none => none
 | some sub => some (rewriteTypedTerm decls subst mono sub)
 .debug (rewriteTyp subst mono τ) e l t' (rewriteTypedTerm decls subst mono r)
 | .retGroup τ e name inner =>
 .retGroup (rewriteTyp subst mono τ) e name (rewriteTypedTerm decls subst mono inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -625,6 +629,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
 let seen := collectInTyp seen τ
 let seen := match t with | some t => collectInTypedTerm seen t | none => seen
 collectInTypedTerm seen r
 | .retGroup τ _ _ inner => collectInTypedTerm (collectInTyp seen τ) inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -683,6 +688,7 @@ def collectCalls (decls : Typed.Decls)
 | .debug _ _ _ t r =>
 let seen := match t with | some t => collectCalls decls seen t | none => seen
 collectCalls decls seen r
 | .retGroup _ _ _ inner => collectCalls decls seen inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -771,6 +777,8 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
 | none => none
 | some sub => some (substInTypedTerm subst sub)
 .debug (Typ.instantiate subst τ) e l t' (substInTypedTerm subst r)
 | .retGroup τ e name inner =>
 .retGroup (Typ.instantiate subst τ) e name (substInTypedTerm subst inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand Down

7 changes: 4 additions & 3 deletions
Ix/Aiur/Compiler/Dedup.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -31,7 +31,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
 (match def_ with | none => none | some b => some (skeletonBlock b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
Expand DownExpand Up@@ -107,7 +107,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
 (match def_ with | none => none | some b => some (rewriteBlock f b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
Expand DownExpand Up@@ -196,7 +196,8 @@ def deduplicate_newFunctions (functions : Array Function) (classes : Array Nat)
 if can then
 let entry := deduplicate_class_entry functions classes cls
 let body := rewriteBlock remapFn f.body
 acc.push { body, layout := f.layout, entry, constrained := false }
 acc.push { body, layout := f.layout, groupNames := f.groupNames,
 entry, constrained := false }
 else acc)
 #[]

Expand Down

74 changes: 61 additions & 13 deletions
Ix/Aiur/Compiler/Lower.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -94,6 +94,15 @@ structure CompilerState where
 ops : Array Bytecode.Op
 selIdx : Bytecode.SelIdx
 degrees : Array Nat
 /-- Top of the `#[return_group(…)]` annotation stack — the display name
 whose index will tag every `Ctrl.return` emitted inside its scope. The
 index is looked up (or allocated) lazily at emit time. -/
 currentReturnGroupName : String := ""
 /-- Group display names allocated so far; position `i` is the name for
 group index `i`. -/
 groupNames : Array String := #[]
 /-- Inverse of `groupNames`: maps name → allocated index. -/
 groupNameMap : Std.HashMap String USize := {}
 deriving Inhabited

abbrev CompileM := EStateM String CompilerState
Expand DownExpand Up@@ -122,6 +131,20 @@ def pushOp (op : Bytecode.Op) (size : Nat := 1) : CompileM (Array Bytecode.ValId
def extractOps : CompileM (Array Bytecode.Op) :=
 modifyGet fun s => (s.ops, {s with ops := #[]})

/-- Look up the `USize` index for the current return-group name, allocating
fresh storage in `groupNames`/`groupNameMap` on first encounter. -/
def allocCurrentGroup : CompileM USize := do
 let st ← get
 let name := st.currentReturnGroupName
 match st.groupNameMap[name]? with
 | some idx => pure idx
 | none =>
 let idx : USize := USize.ofNat st.groupNames.size
 modify fun s => { s with
 groupNameMap := s.groupNameMap.insert name idx
 groupNames := s.groupNames.push name }
 pure idx

open Concrete in
mutual

Expand DownExpand Up@@ -299,6 +322,7 @@ def toIndex
 | some sub => do pure (some (← toIndex layoutMap bindings sub))
 modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
 toIndex layoutMap bindings ret
 | .retGroup _ _ _ inner => toIndex layoutMap bindings inner
termination_by (sizeOf term, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -446,6 +470,12 @@ def Concrete.Term.compile
 let data ← toIndex layoutMap bindings data
 modify fun stt => { stt with ops := stt.ops.push (.ioWrite data) }
 ret.compile returnTyp layoutMap bindings yieldCtrl
 | .retGroup _ _ name inner => do
 let oldGroup := (← get).currentReturnGroupName
 modify fun s => { s with currentReturnGroupName := name }
 let blk ← inner.compile returnTyp layoutMap bindings yieldCtrl
 modify fun s => { s with currentReturnGroupName := oldGroup }
 pure blk
 | .match _ _ scrut cases defaultOpt => do
 let idxs := bindings[scrut]?.getD #[0]
 let ops ← extractOps
Expand All@@ -460,21 +490,24 @@ def Concrete.Term.compile
 pure ({ ops, ctrl } : Bytecode.Block)
 | .ret _ _ term => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 pure ({ ops, ctrl := .return (id - 1) idxs } : Bytecode.Block)
 pure ({ ops, ctrl := .return (id - 1) groupIdx idxs } : Bytecode.Block)
 | _ => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 let ctrl : Bytecode.Ctrl :=
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs else .return (id - 1) idxs
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs
 else .return (id - 1) groupIdx idxs
 pure ({ ops, ctrl } : Bytecode.Block)
termination_by (sizeOf term, 0)
decreasing_by
Expand All@@ -497,7 +530,10 @@ def Concrete.addCase
 | .field g => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (g, term), defaultBlock)
 | .ref global pats => do
 let (index, offsets) ← match layoutMap[global]? with
Expand All@@ -516,22 +552,30 @@ def Concrete.addCase
 acc.insert patLocal slice
 let initState ← get
 let term ← term.compile returnTyp layoutMap ptrBindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (.ofNat index, term), defaultBlock)
 | .wildcard => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases, .some term)
 | _ => throw "addCase: unsupported pattern in concrete lower"
termination_by _ pair => (sizeOf pair.snd, 1)
decreasing_by all_goals first | decreasing_tactic | grind

end

/-- Lower a full concrete function to bytecode. -/
/-- Lower a full concrete function to bytecode. Returns the body, layout
state, and the per-function `groupNames` table (position `i` is the display
name for group index `i`). -/
def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 Except String (Bytecode.Block × Bytecode.LayoutMState) := do
 Except String (Bytecode.Block × Bytecode.LayoutMState × Array String) := do
 let (_inputSize, _outputSize) ← match layoutMap[f.name]? with
 | some (.function layout) => pure (layout.inputSize, layout.outputSize)
 | _ => throw s!"`{f.name}` should be a function"
Expand All@@ -542,15 +586,16 @@ def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 | .ok len => pure len
 let indices := Array.range' valIdx len
 pure (valIdx + len, bindings.insert arg indices)
 let state := { valIdx, selIdx := 0, ops := #[], degrees := Array.replicate valIdx 1 }
 let state : CompilerState := { valIdx, selIdx := 0, ops := #[],
 degrees := Array.replicate valIdx 1 }
 match f.body.compile f.output layoutMap bindings |>.run state with
 | .error e _ => throw e
 | .ok body _ =>
 | .ok body finalState =>
 let (_, layoutMState) := Bytecode.blockLayout body |>.run (.new valIdx)
 let layoutMState := { layoutMState with functionLayout :=
 { layoutMState.functionLayout with
 lookups := layoutMState.functionLayout.lookups + 1 } }
 pure (body, layoutMState)
 pure (body, layoutMState, finalState.groupNames)

def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 Except String (Bytecode.Toplevel × Std.HashMap Global Bytecode.FunIdx) := do
Expand All@@ -559,13 +604,16 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 let (functions, memSizes, nameMap) ← decls.foldlM (init := (#[], initMemSizes, {}))
 fun acc@(functions, memSizes, nameMap) (_, decl) => match decl with
 | .function function => do
 let (body, layoutMState) ← function.compile layout
 let (body, layoutMState, groupNames) ← function.compile layout
 let nameMap := nameMap.insert function.name functions.size
 let function := ⟨body, layoutMState.functionLayout, function.entry, false⟩
 let groupNames := if groupNames.isEmpty then #[""] else groupNames
 let function : Bytecode.Function :=
 { body, layout := layoutMState.functionLayout,
 groupNames, entry := function.entry, constrained := false }
 let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes
 pure (functions.push function, memSizes, nameMap)
 | _ => pure acc
 pure (⟨functions, memSizes.toArray⟩, nameMap)
 pure ({ functions, memorySizes := memSizes.toArray : Bytecode.Toplevel }, nameMap)

end Aiur

Expand Down

1 change: 1 addition & 0 deletions
Ix/Aiur/Compiler/Match.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -389,6 +389,7 @@ def typedToSimple : Term → Simple.Term
 | .debug τ e l t r =>
 let t' := match t with | none => none | some sub => some (typedToSimple sub)
 .debug τ e l t' (typedToSimple r)
 | .retGroup τ e name inner => .retGroup τ e name (typedToSimple inner)
termination_by t => sizeOf t
decreasing_by all_goals first | decreasing_tactic | grind

Expand Down

3 changes: 3 additions & 0 deletions
Ix/Aiur/Compiler/Simple.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -108,6 +108,9 @@ def simplifyTypedTerm (decls : Source.Decls) : Term → Except CheckError Term
 let a' ← simplifyTypedTerm decls a
 let b' ← simplifyTypedTerm decls b
 pure (.u32LessThan τ e a' b')
 | .retGroup τ e name inner => do
 let inner' ← simplifyTypedTerm decls inner
 pure (.retGroup τ e name inner')
 | t => pure t
termination_by t => sizeOf t
decreasing_by
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Remove or un-stick sticky/fixed headers that block content\n(function() {\n function unstick() {\n document.querySelectorAll('header, nav, [role=\"banner\"], .header, .navbar, .sticky, .fixed-top, [style*=\"position: fixed\"], [style*=\"position:sticky\"]').forEach(function(el) {\n if (el.style.position === 'fixed' || el.style.position === 'sticky' || \n getComputedStyle(el).position === 'fixed' || getComputedStyle(el).position === 'sticky') {\n el.style.position = 'static';\n el.style.top = 'auto';\n el.style.zIndex = 'auto';\n }\n });\n }\n \n unstick();\n \n var observer = new MutationObserver(unstick);\n observer.observe(document.body, { childList: true, subtree: true, attributes: true, attributeFilter: ['style', 'class'] });\n})();", "Kill Sticky Headers");
}
} catch(__e) { console.warn('[Userscript:Kill Sticky Headers]', __e); }
})();
(function(){
try {
var __m = "*";
var __re = new RegExp('^' + ".*" + '

Skip to content
Return groups#418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

gabriel-barrett
wants to merge
13
commits into

base:main
Choose a base branch

from

 Draft


Return groups
#418
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


13 commits


Select commit
Hold shift + click to select a range
be60c40

Return groups added to Source


gabriel-barrett May 19, 2026
70dcbfb

Return groups added to other states


gabriel-barrett May 19, 2026
843befa

Added groups to Lean `Bytecode.return`


gabriel-barrett May 19, 2026
a2a7b7f

Added return groups to `bytecode.rs`


gabriel-barrett May 19, 2026
5f7559a

Implemented `Function.split`


gabriel-barrett May 19, 2026
e2528b6

Layout fix on split functions


gabriel-barrett May 19, 2026
3fbc821

Per group circuits


gabriel-barrett May 19, 2026
fefc240

`String` -> `Arc<str>`


gabriel-barrett May 19, 2026
8de0ae2

`address_eq` fast and slow paths


gabriel-barrett May 20, 2026
876e06c

Move split + layout from Rust to Lean


arthurpaulino May 20, 2026
b529b2a

Index return groups by USize at bytecode level


arthurpaulino May 20, 2026
42c4ceb

Use named structures for QueryCounts


arthurpaulino May 20, 2026
60c38d4

Preserve groupNames through dedup + match arms


arthurpaulino May 20, 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


1 change: 1 addition & 0 deletions
Ix/Aiur.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,5 +20,6 @@ public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Layout
public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Split
public import Ix.Aiur.Compiler
public import Ix.Aiur.Statistics

6 changes: 4 additions & 2 deletions
Ix/Aiur/Compiler.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@ public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Simple
public import Ix.Aiur.Compiler.Split

/-!
Aiur compiler pipeline: type-check, simplify, concretize, lower, deduplicate.
Expand DownExpand Up@@ -68,7 +69,7 @@ def Bytecode.Ctrl.collectConstrainedCallees (c : Bytecode.Ctrl) :
 | some block => branchCallees ++ block.collectConstrainedCallees
 | none => branchCallees
 withDefault ++ continuation.collectConstrainedCallees
 | .return _ _ | .yield _ _ => #[]
 | .return _ _ _ | .yield _ _ => #[]
termination_by (sizeOf c, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -117,9 +118,10 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev
 let (bytecodeRaw, preNameMap) ← concDecls.toBytecode
 let (bytecodeDedup, remap) := bytecodeRaw.deduplicate
 let needs := bytecodeDedup.needsCircuit
 let bytecode := { bytecodeDedup with
 let bytecodeConstrained : Bytecode.Toplevel := { bytecodeDedup with
 functions := bytecodeDedup.functions.mapIdx fun i f =>
 { f with constrained := needs[i]! } }
 let bytecode := bytecodeConstrained.computeFiltered
 let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx))
 fun acc name idx => acc.insert name (remap idx)
 pure (CompiledToplevel.mk t bytecode nameMap)
Expand Down

5 changes: 5 additions & 0 deletions
Ix/Aiur/Compiler/Check.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -790,6 +790,9 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
 | some sub => do pure (some (← inferTerm sub))
 let ret' ← inferTerm ret
 pure (Typed.Term.debug ret'.typ ret'.escapes label term' ret')
 | .retGroup name inner => do
 let inner' ← inferTerm inner
 pure (Typed.Term.retGroup inner'.typ inner'.escapes name inner')
termination_by (sizeOf t, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -917,6 +920,8 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
 | none => pure none
 | some sub => do pure (some (← zonkTypedTerm sub))
 pure (.debug (← zonkTyp τ) e label t' (← zonkTypedTerm r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← zonkTyp τ) e name (← zonkTypedTerm inner))
termination_by sizeOf t
decreasing_by
 all_goals first
Expand Down

8 changes: 8 additions & 0 deletions
Ix/Aiur/Compiler/Concretize.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -349,6 +349,8 @@ def termToConcrete
 | none => pure none
 | some sub => do pure (some (← termToConcrete mono sub))
 pure (.debug (← typToConcrete mono τ) e l t' (← termToConcrete mono r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← typToConcrete mono τ) e name (← termToConcrete mono inner))
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -541,6 +543,8 @@ def rewriteTypedTerm (decls : Typed.Decls)
 | none => none
 | some sub => some (rewriteTypedTerm decls subst mono sub)
 .debug (rewriteTyp subst mono τ) e l t' (rewriteTypedTerm decls subst mono r)
 | .retGroup τ e name inner =>
 .retGroup (rewriteTyp subst mono τ) e name (rewriteTypedTerm decls subst mono inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -625,6 +629,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
 let seen := collectInTyp seen τ
 let seen := match t with | some t => collectInTypedTerm seen t | none => seen
 collectInTypedTerm seen r
 | .retGroup τ _ _ inner => collectInTypedTerm (collectInTyp seen τ) inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -683,6 +688,7 @@ def collectCalls (decls : Typed.Decls)
 | .debug _ _ _ t r =>
 let seen := match t with | some t => collectCalls decls seen t | none => seen
 collectCalls decls seen r
 | .retGroup _ _ _ inner => collectCalls decls seen inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -771,6 +777,8 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
 | none => none
 | some sub => some (substInTypedTerm subst sub)
 .debug (Typ.instantiate subst τ) e l t' (substInTypedTerm subst r)
 | .retGroup τ e name inner =>
 .retGroup (Typ.instantiate subst τ) e name (substInTypedTerm subst inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand Down

7 changes: 4 additions & 3 deletions
Ix/Aiur/Compiler/Dedup.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -31,7 +31,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
 (match def_ with | none => none | some b => some (skeletonBlock b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
Expand DownExpand Up@@ -107,7 +107,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
 (match def_ with | none => none | some b => some (rewriteBlock f b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
Expand DownExpand Up@@ -196,7 +196,8 @@ def deduplicate_newFunctions (functions : Array Function) (classes : Array Nat)
 if can then
 let entry := deduplicate_class_entry functions classes cls
 let body := rewriteBlock remapFn f.body
 acc.push { body, layout := f.layout, entry, constrained := false }
 acc.push { body, layout := f.layout, groupNames := f.groupNames,
 entry, constrained := false }
 else acc)
 #[]

Expand Down

74 changes: 61 additions & 13 deletions
Ix/Aiur/Compiler/Lower.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -94,6 +94,15 @@ structure CompilerState where
 ops : Array Bytecode.Op
 selIdx : Bytecode.SelIdx
 degrees : Array Nat
 /-- Top of the `#[return_group(…)]` annotation stack — the display name
 whose index will tag every `Ctrl.return` emitted inside its scope. The
 index is looked up (or allocated) lazily at emit time. -/
 currentReturnGroupName : String := ""
 /-- Group display names allocated so far; position `i` is the name for
 group index `i`. -/
 groupNames : Array String := #[]
 /-- Inverse of `groupNames`: maps name → allocated index. -/
 groupNameMap : Std.HashMap String USize := {}
 deriving Inhabited

abbrev CompileM := EStateM String CompilerState
Expand DownExpand Up@@ -122,6 +131,20 @@ def pushOp (op : Bytecode.Op) (size : Nat := 1) : CompileM (Array Bytecode.ValId
def extractOps : CompileM (Array Bytecode.Op) :=
 modifyGet fun s => (s.ops, {s with ops := #[]})

/-- Look up the `USize` index for the current return-group name, allocating
fresh storage in `groupNames`/`groupNameMap` on first encounter. -/
def allocCurrentGroup : CompileM USize := do
 let st ← get
 let name := st.currentReturnGroupName
 match st.groupNameMap[name]? with
 | some idx => pure idx
 | none =>
 let idx : USize := USize.ofNat st.groupNames.size
 modify fun s => { s with
 groupNameMap := s.groupNameMap.insert name idx
 groupNames := s.groupNames.push name }
 pure idx

open Concrete in
mutual

Expand DownExpand Up@@ -299,6 +322,7 @@ def toIndex
 | some sub => do pure (some (← toIndex layoutMap bindings sub))
 modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
 toIndex layoutMap bindings ret
 | .retGroup _ _ _ inner => toIndex layoutMap bindings inner
termination_by (sizeOf term, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -446,6 +470,12 @@ def Concrete.Term.compile
 let data ← toIndex layoutMap bindings data
 modify fun stt => { stt with ops := stt.ops.push (.ioWrite data) }
 ret.compile returnTyp layoutMap bindings yieldCtrl
 | .retGroup _ _ name inner => do
 let oldGroup := (← get).currentReturnGroupName
 modify fun s => { s with currentReturnGroupName := name }
 let blk ← inner.compile returnTyp layoutMap bindings yieldCtrl
 modify fun s => { s with currentReturnGroupName := oldGroup }
 pure blk
 | .match _ _ scrut cases defaultOpt => do
 let idxs := bindings[scrut]?.getD #[0]
 let ops ← extractOps
Expand All@@ -460,21 +490,24 @@ def Concrete.Term.compile
 pure ({ ops, ctrl } : Bytecode.Block)
 | .ret _ _ term => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 pure ({ ops, ctrl := .return (id - 1) idxs } : Bytecode.Block)
 pure ({ ops, ctrl := .return (id - 1) groupIdx idxs } : Bytecode.Block)
 | _ => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 let ctrl : Bytecode.Ctrl :=
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs else .return (id - 1) idxs
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs
 else .return (id - 1) groupIdx idxs
 pure ({ ops, ctrl } : Bytecode.Block)
termination_by (sizeOf term, 0)
decreasing_by
Expand All@@ -497,7 +530,10 @@ def Concrete.addCase
 | .field g => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (g, term), defaultBlock)
 | .ref global pats => do
 let (index, offsets) ← match layoutMap[global]? with
Expand All@@ -516,22 +552,30 @@ def Concrete.addCase
 acc.insert patLocal slice
 let initState ← get
 let term ← term.compile returnTyp layoutMap ptrBindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (.ofNat index, term), defaultBlock)
 | .wildcard => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases, .some term)
 | _ => throw "addCase: unsupported pattern in concrete lower"
termination_by _ pair => (sizeOf pair.snd, 1)
decreasing_by all_goals first | decreasing_tactic | grind

end

/-- Lower a full concrete function to bytecode. -/
/-- Lower a full concrete function to bytecode. Returns the body, layout
state, and the per-function `groupNames` table (position `i` is the display
name for group index `i`). -/
def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 Except String (Bytecode.Block × Bytecode.LayoutMState) := do
 Except String (Bytecode.Block × Bytecode.LayoutMState × Array String) := do
 let (_inputSize, _outputSize) ← match layoutMap[f.name]? with
 | some (.function layout) => pure (layout.inputSize, layout.outputSize)
 | _ => throw s!"`{f.name}` should be a function"
Expand All@@ -542,15 +586,16 @@ def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 | .ok len => pure len
 let indices := Array.range' valIdx len
 pure (valIdx + len, bindings.insert arg indices)
 let state := { valIdx, selIdx := 0, ops := #[], degrees := Array.replicate valIdx 1 }
 let state : CompilerState := { valIdx, selIdx := 0, ops := #[],
 degrees := Array.replicate valIdx 1 }
 match f.body.compile f.output layoutMap bindings |>.run state with
 | .error e _ => throw e
 | .ok body _ =>
 | .ok body finalState =>
 let (_, layoutMState) := Bytecode.blockLayout body |>.run (.new valIdx)
 let layoutMState := { layoutMState with functionLayout :=
 { layoutMState.functionLayout with
 lookups := layoutMState.functionLayout.lookups + 1 } }
 pure (body, layoutMState)
 pure (body, layoutMState, finalState.groupNames)

def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 Except String (Bytecode.Toplevel × Std.HashMap Global Bytecode.FunIdx) := do
Expand All@@ -559,13 +604,16 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 let (functions, memSizes, nameMap) ← decls.foldlM (init := (#[], initMemSizes, {}))
 fun acc@(functions, memSizes, nameMap) (_, decl) => match decl with
 | .function function => do
 let (body, layoutMState) ← function.compile layout
 let (body, layoutMState, groupNames) ← function.compile layout
 let nameMap := nameMap.insert function.name functions.size
 let function := ⟨body, layoutMState.functionLayout, function.entry, false⟩
 let groupNames := if groupNames.isEmpty then #[""] else groupNames
 let function : Bytecode.Function :=
 { body, layout := layoutMState.functionLayout,
 groupNames, entry := function.entry, constrained := false }
 let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes
 pure (functions.push function, memSizes, nameMap)
 | _ => pure acc
 pure (⟨functions, memSizes.toArray⟩, nameMap)
 pure ({ functions, memorySizes := memSizes.toArray : Bytecode.Toplevel }, nameMap)

end Aiur

Expand Down

1 change: 1 addition & 0 deletions
Ix/Aiur/Compiler/Match.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -389,6 +389,7 @@ def typedToSimple : Term → Simple.Term
 | .debug τ e l t r =>
 let t' := match t with | none => none | some sub => some (typedToSimple sub)
 .debug τ e l t' (typedToSimple r)
 | .retGroup τ e name inner => .retGroup τ e name (typedToSimple inner)
termination_by t => sizeOf t
decreasing_by all_goals first | decreasing_tactic | grind

Expand Down

3 changes: 3 additions & 0 deletions
Ix/Aiur/Compiler/Simple.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -108,6 +108,9 @@ def simplifyTypedTerm (decls : Source.Decls) : Term → Except CheckError Term
 let a' ← simplifyTypedTerm decls a
 let b' ← simplifyTypedTerm decls b
 pure (.u32LessThan τ e a' b')
 | .retGroup τ e name inner => do
 let inner' ← simplifyTypedTerm decls inner
 pure (.retGroup τ e name inner')
 | t => pure t
termination_by t => sizeOf t
decreasing_by
Expand Down
Loading

, 'i');
if (__m === '*' || __re.test(location.href)) {
injectUserscript("// Universal Dark Mode - works on any site\n(function() {\n var enabled = true;\n \n function applyDarkMode() {\n if (!enabled) return;\n \n // Create style element if it doesn't exist\n var style = document.getElementById('universal-dark-mode-style');\n if (!style) {\n style = document.createElement('style');\n style.id = 'universal-dark-mode-style';\n document.head.appendChild(style);\n }\n \n // Dark mode CSS - inverts colors but preserves images/video\n style.textContent = '\n /* Invert everything except media */\n html {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #1a1a2e !important;\n }\n \n /* Restore images, videos, iframes, canvas */\n img, video, iframe, canvas, svg, picture, [style*=\"background-image\"] {\n filter: invert(1) hue-rotate(180deg) !important;\n }\n \n /* Preserve specific elements that should not be inverted */\n .no-dark-mode, .no-dark-mode *,\n [data-theme=\"light\"], [data-theme=\"light\"],\n .ace_editor, .ace_editor *,\n .CodeMirror, .CodeMirror *,\n .monaco-editor, .monaco-editor *,\n .markdown-body pre, .markdown-body pre *,\n .highlight, .highlight *,\n pre code, pre code * {\n filter: none !important;\n }\n \n /* Fix common UI elements */\n .modal, .popup, .dropdown-menu, .tooltip, .popover {\n filter: invert(1) hue-rotate(180deg) !important;\n background: #2d2d44 !important;\n border-color: #444 !important;\n }\n \n /* Scrollbars */\n ::-webkit-scrollbar { background: #1a1a2e !important; }\n ::-webkit-scrollbar-thumb { background: #444 !important; }\n ::-webkit-scrollbar-thumb:hover { background: #555 !important; }\n \n /* Selection */\n ::selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ::-moz-selection { background: #4ecdc4 !important; color: #1a1a2e !important; }\n ';\n }\n \n function removeDarkMode() {\n var style = document.getElementById('universal-dark-mode-style');\n if (style) style.remove();\n }\n \n // Toggle with Alt+Shift+D\n document.addEventListener('keydown', function(e) {\n if (e.altKey && e.shiftKey && e.key === 'D') {\n e.preventDefault();\n enabled = !enabled;\n if (enabled) {\n applyDarkMode();\n console.log('[Universal Dark Mode] Enabled');\n } else {\n removeDarkMode();\n console.log('[Universal Dark Mode] Disabled');\n }\n }\n });\n \n // Apply on load\n applyDarkMode();\n \n // Re-apply on dynamic content\n var observer = new MutationObserver(function(mutations) {\n if (enabled && !document.getElementById('universal-dark-mode-style')) {\n applyDarkMode();\n }\n });\n observer.observe(document.head, { childList: true });\n \n console.log('[Universal Dark Mode] Loaded - Press Alt+Shift+D to toggle');\n})();", "Universal Dark Mode");
}
} catch(__e) { console.warn('[Userscript:Universal Dark Mode]', __e); }
})();
})();

Skip to content
Return groups#418

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

 Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and
privacy statement. We’ll occasionally send you account related emails.

Already on GitHub?
Sign in
to your account

 Draft

gabriel-barrett
wants to merge
13
commits into

base:main
Choose a base branch

from

 Draft


Return groups
#418
Show file tree
Hide file tree
 Changes from all commits
Commits

Show all changes


13 commits


Select commit
Hold shift + click to select a range
be60c40

Return groups added to Source


gabriel-barrett May 19, 2026
70dcbfb

Return groups added to other states


gabriel-barrett May 19, 2026
843befa

Added groups to Lean `Bytecode.return`


gabriel-barrett May 19, 2026
a2a7b7f

Added return groups to `bytecode.rs`


gabriel-barrett May 19, 2026
5f7559a

Implemented `Function.split`


gabriel-barrett May 19, 2026
e2528b6

Layout fix on split functions


gabriel-barrett May 19, 2026
3fbc821

Per group circuits


gabriel-barrett May 19, 2026
fefc240

`String` -> `Arc<str>`


gabriel-barrett May 19, 2026
8de0ae2

`address_eq` fast and slow paths


gabriel-barrett May 20, 2026
876e06c

Move split + layout from Rust to Lean


arthurpaulino May 20, 2026
b529b2a

Index return groups by USize at bytecode level


arthurpaulino May 20, 2026
42c4ceb

Use named structures for QueryCounts


arthurpaulino May 20, 2026
60c38d4

Preserve groupNames through dedup + match arms


arthurpaulino May 20, 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


1 change: 1 addition & 0 deletions
Ix/Aiur.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -20,5 +20,6 @@ public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Layout
public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Split
public import Ix.Aiur.Compiler
public import Ix.Aiur.Statistics

6 changes: 4 additions & 2 deletions
Ix/Aiur/Compiler.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -3,6 +3,7 @@ public import Ix.Aiur.Compiler.Lower
public import Ix.Aiur.Compiler.Dedup
public import Ix.Aiur.Compiler.Concretize
public import Ix.Aiur.Compiler.Simple
public import Ix.Aiur.Compiler.Split

/-!
Aiur compiler pipeline: type-check, simplify, concretize, lower, deduplicate.
Expand DownExpand Up@@ -68,7 +69,7 @@ def Bytecode.Ctrl.collectConstrainedCallees (c : Bytecode.Ctrl) :
 | some block => branchCallees ++ block.collectConstrainedCallees
 | none => branchCallees
 withDefault ++ continuation.collectConstrainedCallees
 | .return _ _ | .yield _ _ => #[]
 | .return _ _ _ | .yield _ _ => #[]
termination_by (sizeOf c, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -117,9 +118,10 @@ def Source.Toplevel.compile (t : Source.Toplevel) : Except String CompiledToplev
 let (bytecodeRaw, preNameMap) ← concDecls.toBytecode
 let (bytecodeDedup, remap) := bytecodeRaw.deduplicate
 let needs := bytecodeDedup.needsCircuit
 let bytecode := { bytecodeDedup with
 let bytecodeConstrained : Bytecode.Toplevel := { bytecodeDedup with
 functions := bytecodeDedup.functions.mapIdx fun i f =>
 { f with constrained := needs[i]! } }
 let bytecode := bytecodeConstrained.computeFiltered
 let nameMap := preNameMap.fold (init := (∅ : Std.HashMap Global Bytecode.FunIdx))
 fun acc name idx => acc.insert name (remap idx)
 pure (CompiledToplevel.mk t bytecode nameMap)
Expand Down

5 changes: 5 additions & 0 deletions
Ix/Aiur/Compiler/Check.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -790,6 +790,9 @@ def inferTerm (t : Term) : CheckM Typed.Term := match t with
 | some sub => do pure (some (← inferTerm sub))
 let ret' ← inferTerm ret
 pure (Typed.Term.debug ret'.typ ret'.escapes label term' ret')
 | .retGroup name inner => do
 let inner' ← inferTerm inner
 pure (Typed.Term.retGroup inner'.typ inner'.escapes name inner')
termination_by (sizeOf t, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -917,6 +920,8 @@ def zonkTypedTerm (t : Typed.Term) : CheckM Typed.Term := match t with
 | none => pure none
 | some sub => do pure (some (← zonkTypedTerm sub))
 pure (.debug (← zonkTyp τ) e label t' (← zonkTypedTerm r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← zonkTyp τ) e name (← zonkTypedTerm inner))
termination_by sizeOf t
decreasing_by
 all_goals first
Expand Down

8 changes: 8 additions & 0 deletions
Ix/Aiur/Compiler/Concretize.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -349,6 +349,8 @@ def termToConcrete
 | none => pure none
 | some sub => do pure (some (← termToConcrete mono sub))
 pure (.debug (← typToConcrete mono τ) e l t' (← termToConcrete mono r))
 | .retGroup τ e name inner => do
 pure (.retGroup (← typToConcrete mono τ) e name (← termToConcrete mono inner))
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -541,6 +543,8 @@ def rewriteTypedTerm (decls : Typed.Decls)
 | none => none
 | some sub => some (rewriteTypedTerm decls subst mono sub)
 .debug (rewriteTyp subst mono τ) e l t' (rewriteTypedTerm decls subst mono r)
 | .retGroup τ e name inner =>
 .retGroup (rewriteTyp subst mono τ) e name (rewriteTypedTerm decls subst mono inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -625,6 +629,7 @@ def collectInTypedTerm (seen : Std.HashSet (Global × Array Typ)) :
 let seen := collectInTyp seen τ
 let seen := match t with | some t => collectInTypedTerm seen t | none => seen
 collectInTypedTerm seen r
 | .retGroup τ _ _ inner => collectInTypedTerm (collectInTyp seen τ) inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -683,6 +688,7 @@ def collectCalls (decls : Typed.Decls)
 | .debug _ _ _ t r =>
 let seen := match t with | some t => collectCalls decls seen t | none => seen
 collectCalls decls seen r
 | .retGroup _ _ _ inner => collectCalls decls seen inner
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand DownExpand Up@@ -771,6 +777,8 @@ def substInTypedTerm (subst : Global → Option Typ) : Typed.Term → Typed.Term
 | none => none
 | some sub => some (substInTypedTerm subst sub)
 .debug (Typ.instantiate subst τ) e l t' (substInTypedTerm subst r)
 | .retGroup τ e name inner =>
 .retGroup (Typ.instantiate subst τ) e name (substInTypedTerm subst inner)
termination_by t => sizeOf t
decreasing_by
 all_goals first
Expand Down

7 changes: 4 additions & 3 deletions
Ix/Aiur/Compiler/Dedup.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -31,7 +31,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
 (match def_ with | none => none | some b => some (skeletonBlock b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, skeletonBlock b))
Expand DownExpand Up@@ -107,7 +107,7 @@ mutual
 | .match v branches def_ =>
 .match v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
 (match def_ with | none => none | some b => some (rewriteBlock f b))
 | .return s vs => .return s vs
 | .return s g vs => .return s g vs
 | .yield s vs => .yield s vs
 | .matchContinue v branches def_ outputSize sharedAux sharedLookups cont =>
 .matchContinue v (branches.attach.map fun ⟨(g, b), _⟩ => (g, rewriteBlock f b))
Expand DownExpand Up@@ -196,7 +196,8 @@ def deduplicate_newFunctions (functions : Array Function) (classes : Array Nat)
 if can then
 let entry := deduplicate_class_entry functions classes cls
 let body := rewriteBlock remapFn f.body
 acc.push { body, layout := f.layout, entry, constrained := false }
 acc.push { body, layout := f.layout, groupNames := f.groupNames,
 entry, constrained := false }
 else acc)
 #[]

Expand Down

74 changes: 61 additions & 13 deletions
Ix/Aiur/Compiler/Lower.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -94,6 +94,15 @@ structure CompilerState where
 ops : Array Bytecode.Op
 selIdx : Bytecode.SelIdx
 degrees : Array Nat
 /-- Top of the `#[return_group(…)]` annotation stack — the display name
 whose index will tag every `Ctrl.return` emitted inside its scope. The
 index is looked up (or allocated) lazily at emit time. -/
 currentReturnGroupName : String := ""
 /-- Group display names allocated so far; position `i` is the name for
 group index `i`. -/
 groupNames : Array String := #[]
 /-- Inverse of `groupNames`: maps name → allocated index. -/
 groupNameMap : Std.HashMap String USize := {}
 deriving Inhabited

abbrev CompileM := EStateM String CompilerState
Expand DownExpand Up@@ -122,6 +131,20 @@ def pushOp (op : Bytecode.Op) (size : Nat := 1) : CompileM (Array Bytecode.ValId
def extractOps : CompileM (Array Bytecode.Op) :=
 modifyGet fun s => (s.ops, {s with ops := #[]})

/-- Look up the `USize` index for the current return-group name, allocating
fresh storage in `groupNames`/`groupNameMap` on first encounter. -/
def allocCurrentGroup : CompileM USize := do
 let st ← get
 let name := st.currentReturnGroupName
 match st.groupNameMap[name]? with
 | some idx => pure idx
 | none =>
 let idx : USize := USize.ofNat st.groupNames.size
 modify fun s => { s with
 groupNameMap := s.groupNameMap.insert name idx
 groupNames := s.groupNames.push name }
 pure idx

open Concrete in
mutual

Expand DownExpand Up@@ -299,6 +322,7 @@ def toIndex
 | some sub => do pure (some (← toIndex layoutMap bindings sub))
 modify fun stt => { stt with ops := stt.ops.push (.debug label term) }
 toIndex layoutMap bindings ret
 | .retGroup _ _ _ inner => toIndex layoutMap bindings inner
termination_by (sizeOf term, 0)
decreasing_by
 all_goals first
Expand DownExpand Up@@ -446,6 +470,12 @@ def Concrete.Term.compile
 let data ← toIndex layoutMap bindings data
 modify fun stt => { stt with ops := stt.ops.push (.ioWrite data) }
 ret.compile returnTyp layoutMap bindings yieldCtrl
 | .retGroup _ _ name inner => do
 let oldGroup := (← get).currentReturnGroupName
 modify fun s => { s with currentReturnGroupName := name }
 let blk ← inner.compile returnTyp layoutMap bindings yieldCtrl
 modify fun s => { s with currentReturnGroupName := oldGroup }
 pure blk
 | .match _ _ scrut cases defaultOpt => do
 let idxs := bindings[scrut]?.getD #[0]
 let ops ← extractOps
Expand All@@ -460,21 +490,24 @@ def Concrete.Term.compile
 pure ({ ops, ctrl } : Bytecode.Block)
 | .ret _ _ term => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 pure ({ ops, ctrl := .return (id - 1) idxs } : Bytecode.Block)
 pure ({ ops, ctrl := .return (id - 1) groupIdx idxs } : Bytecode.Block)
 | _ => do
 let idxs ← toIndex layoutMap bindings term
 let groupIdx ← allocCurrentGroup
 let state ← get
 let state := { state with selIdx := state.selIdx + 1 }
 set state
 let ops := state.ops
 let id := state.selIdx
 let ctrl : Bytecode.Ctrl :=
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs else .return (id - 1) idxs
 if yieldCtrl && !term.escapes then .yield (id - 1) idxs
 else .return (id - 1) groupIdx idxs
 pure ({ ops, ctrl } : Bytecode.Block)
termination_by (sizeOf term, 0)
decreasing_by
Expand All@@ -497,7 +530,10 @@ def Concrete.addCase
 | .field g => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (g, term), defaultBlock)
 | .ref global pats => do
 let (index, offsets) ← match layoutMap[global]? with
Expand All@@ -516,22 +552,30 @@ def Concrete.addCase
 acc.insert patLocal slice
 let initState ← get
 let term ← term.compile returnTyp layoutMap ptrBindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases.push (.ofNat index, term), defaultBlock)
 | .wildcard => do
 let initState ← get
 let term ← term.compile returnTyp layoutMap bindings yieldCtrl
 set { initState with selIdx := (← get).selIdx }
 let cur ← get
 set { initState with selIdx := cur.selIdx,
 groupNames := cur.groupNames,
 groupNameMap := cur.groupNameMap }
 pure (cases, .some term)
 | _ => throw "addCase: unsupported pattern in concrete lower"
termination_by _ pair => (sizeOf pair.snd, 1)
decreasing_by all_goals first | decreasing_tactic | grind

end

/-- Lower a full concrete function to bytecode. -/
/-- Lower a full concrete function to bytecode. Returns the body, layout
state, and the per-function `groupNames` table (position `i` is the display
name for group index `i`). -/
def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 Except String (Bytecode.Block × Bytecode.LayoutMState) := do
 Except String (Bytecode.Block × Bytecode.LayoutMState × Array String) := do
 let (_inputSize, _outputSize) ← match layoutMap[f.name]? with
 | some (.function layout) => pure (layout.inputSize, layout.outputSize)
 | _ => throw s!"`{f.name}` should be a function"
Expand All@@ -542,15 +586,16 @@ def Concrete.Function.compile (layoutMap : LayoutMap) (f : Concrete.Function) :
 | .ok len => pure len
 let indices := Array.range' valIdx len
 pure (valIdx + len, bindings.insert arg indices)
 let state := { valIdx, selIdx := 0, ops := #[], degrees := Array.replicate valIdx 1 }
 let state : CompilerState := { valIdx, selIdx := 0, ops := #[],
 degrees := Array.replicate valIdx 1 }
 match f.body.compile f.output layoutMap bindings |>.run state with
 | .error e _ => throw e
 | .ok body _ =>
 | .ok body finalState =>
 let (_, layoutMState) := Bytecode.blockLayout body |>.run (.new valIdx)
 let layoutMState := { layoutMState with functionLayout :=
 { layoutMState.functionLayout with
 lookups := layoutMState.functionLayout.lookups + 1 } }
 pure (body, layoutMState)
 pure (body, layoutMState, finalState.groupNames)

def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 Except String (Bytecode.Toplevel × Std.HashMap Global Bytecode.FunIdx) := do
Expand All@@ -559,13 +604,16 @@ def Concrete.Decls.toBytecode (decls : Concrete.Decls) :
 let (functions, memSizes, nameMap) ← decls.foldlM (init := (#[], initMemSizes, {}))
 fun acc@(functions, memSizes, nameMap) (_, decl) => match decl with
 | .function function => do
 let (body, layoutMState) ← function.compile layout
 let (body, layoutMState, groupNames) ← function.compile layout
 let nameMap := nameMap.insert function.name functions.size
 let function := ⟨body, layoutMState.functionLayout, function.entry, false⟩
 let groupNames := if groupNames.isEmpty then #[""] else groupNames
 let function : Bytecode.Function :=
 { body, layout := layoutMState.functionLayout,
 groupNames, entry := function.entry, constrained := false }
 let memSizes := layoutMState.memSizes.fold (·.insert ·) memSizes
 pure (functions.push function, memSizes, nameMap)
 | _ => pure acc
 pure (⟨functions, memSizes.toArray⟩, nameMap)
 pure ({ functions, memorySizes := memSizes.toArray : Bytecode.Toplevel }, nameMap)

end Aiur

Expand Down

1 change: 1 addition & 0 deletions
Ix/Aiur/Compiler/Match.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -389,6 +389,7 @@ def typedToSimple : Term → Simple.Term
 | .debug τ e l t r =>
 let t' := match t with | none => none | some sub => some (typedToSimple sub)
 .debug τ e l t' (typedToSimple r)
 | .retGroup τ e name inner => .retGroup τ e name (typedToSimple inner)
termination_by t => sizeOf t
decreasing_by all_goals first | decreasing_tactic | grind

Expand Down

3 changes: 3 additions & 0 deletions
Ix/Aiur/Compiler/Simple.lean

View file


Open in desktop

Original file line numberDiff line numberDiff line change
Expand Up@@ -108,6 +108,9 @@ def simplifyTypedTerm (decls : Source.Decls) : Term → Except CheckError Term
 let a' ← simplifyTypedTerm decls a
 let b' ← simplifyTypedTerm decls b
 pure (.u32LessThan τ e a' b')
 | .retGroup τ e name inner => do
 let inner' ← simplifyTypedTerm decls inner
 pure (.retGroup τ e name inner')
 | t => pure t
termination_by t => sizeOf t
decreasing_by
Expand Down
Loading