@@ -17,6 +17,153 @@ use crate::compiler::{
1717
1818use super :: { BranchState , compile_expr} ;
1919
20+ /// Feature #079 — K-table lookup for dense small-domain numeric `when`.
21+ /// See docs/state/features/079_*.md.
22+ fn const_fold_i64 ( expr : & Expr < ' _ > ) -> Option < i64 > {
23+ match expr {
24+ Expr :: Num ( s, _) => lake_frontend:: api:: expr:: parse_int_literal ( s) . ok ( ) ,
25+ Expr :: Bool ( b) => Some ( if * b { 1 } else { 0 } ) ,
26+ Expr :: Neg ( inner) => const_fold_i64 ( & inner. inner ) . map ( |v| v. wrapping_neg ( ) ) ,
27+ Expr :: Add ( l, r) => Some ( const_fold_i64 ( & l. inner ) ?. wrapping_add ( const_fold_i64 ( & r. inner ) ?) ) ,
28+ Expr :: Sub ( l, r) => Some ( const_fold_i64 ( & l. inner ) ?. wrapping_sub ( const_fold_i64 ( & r. inner ) ?) ) ,
29+ Expr :: Mul ( l, r) => Some ( const_fold_i64 ( & l. inner ) ?. wrapping_mul ( const_fold_i64 ( & r. inner ) ?) ) ,
30+ Expr :: Div ( l, r) => {
31+ let rv = const_fold_i64 ( & r. inner ) ?;
32+ if rv == 0 {
33+ None
34+ } else {
35+ Some ( const_fold_i64 ( & l. inner ) ?. wrapping_div ( rv) )
36+ }
37+ }
38+ Expr :: BAnd ( l, r) => Some ( const_fold_i64 ( & l. inner ) ? & const_fold_i64 ( & r. inner ) ?) ,
39+ Expr :: BOr ( l, r) => Some ( const_fold_i64 ( & l. inner ) ? | const_fold_i64 ( & r. inner ) ?) ,
40+ Expr :: BXor ( l, r) => Some ( const_fold_i64 ( & l. inner ) ? ^ const_fold_i64 ( & r. inner ) ?) ,
41+ Expr :: Shl ( l, r) => {
42+ let rv = const_fold_i64 ( & r. inner ) ?;
43+ if !( 0 ..64 ) . contains ( & rv) {
44+ None
45+ } else {
46+ Some ( const_fold_i64 ( & l. inner ) ?. wrapping_shl ( rv as u32 ) )
47+ }
48+ }
49+ Expr :: Shr ( l, r) => {
50+ let rv = const_fold_i64 ( & r. inner ) ?;
51+ if !( 0 ..64 ) . contains ( & rv) {
52+ None
53+ } else {
54+ Some ( ( const_fold_i64 ( & l. inner ) ? as u64 ) . wrapping_shr ( rv as u32 ) as i64 )
55+ }
56+ }
57+ _ => None ,
58+ }
59+ }
60+
61+ /// Detector: returns `Some(values)` iff every arm key is an i64 `Num`,
62+ /// keys form exactly `0..N`, no wildcard, and every arm body is one
63+ /// const-foldable i64 expression.
64+ fn try_build_k_table < ' a > ( branches : & [ ( Expr < ' a > , Vec < Expr < ' a > > ) ] ) -> Option < Vec < i64 > > {
65+ if branches. len ( ) < 2 {
66+ return None ;
67+ }
68+ let n = branches. len ( ) as i64 ;
69+ let mut values = vec ! [ 0i64 ; branches. len( ) ] ;
70+ let mut seen = vec ! [ false ; branches. len( ) ] ;
71+ for ( cond, body) in branches. iter ( ) {
72+ if is_wildcard ( cond) {
73+ return None ;
74+ }
75+ let key = match cond {
76+ Expr :: Num ( s, _) => lake_frontend:: api:: expr:: parse_int_literal ( s) . ok ( ) ?,
77+ _ => return None ,
78+ } ;
79+ if !( 0 ..n) . contains ( & key) {
80+ return None ;
81+ }
82+ let idx = key as usize ;
83+ if seen[ idx] {
84+ return None ;
85+ }
86+ seen[ idx] = true ;
87+ if body. len ( ) != 1 {
88+ return None ;
89+ }
90+ values[ idx] = const_fold_i64 ( & body[ 0 ] ) ?;
91+ }
92+ if !seen. iter ( ) . all ( |& s| s) {
93+ return None ;
94+ }
95+ Some ( values)
96+ }
97+
98+ /// Emit the K-table lookup path. Assumes the discriminant has already
99+ /// been compiled and lands at `disc_done_id`. Returns the `after_when_id`
100+ /// to use as the continuation point.
101+ fn emit_k_table (
102+ ctx : & mut CompilerCtx ,
103+ builder : & mut FunctionBuilder ,
104+ machine_ctx_var : Variable ,
105+ outer_switch : & mut Switch ,
106+ disc_done_id : i64 ,
107+ values : & [ i64 ] ,
108+ ) -> Result < i64 > {
109+ let ptr_ty = ctx. module ( ) . target_config ( ) . pointer_type ( ) ;
110+ let after_when_id = disc_done_id + 1 ;
111+ let n = values. len ( ) as i64 ;
112+
113+ // Declare a Local .rodata symbol holding the i64 table.
114+ let data_name = format ! ( "ktbl_{disc_done_id}" ) ;
115+ let data_id = ctx
116+ . module_mut ( )
117+ . declare_data ( & data_name, Linkage :: Local , false , false ) ?;
118+ let mut bytes: Vec < u8 > = Vec :: with_capacity ( values. len ( ) * 8 ) ;
119+ for v in values {
120+ bytes. extend_from_slice ( & v. to_le_bytes ( ) ) ;
121+ }
122+ let mut desc = DataDescription :: new ( ) ;
123+ desc. define ( bytes. into_boxed_slice ( ) ) ;
124+ ctx. module_mut ( ) . define_data ( data_id, & desc) ?;
125+
126+ let b_lookup = builder. create_block ( ) ;
127+ let b_inrange = builder. create_block ( ) ;
128+ let b_oob = builder. create_block ( ) ;
129+
130+ // Lookup block — load disc from TEMP_VAL, bounds-check, branch.
131+ // Block sealing is deferred to `seal_all_blocks()` at the end of
132+ // machine compilation: sealing locally with zero predecessors
133+ // (the branch_switch jumps to b_lookup get added later) would
134+ // make `use_var(machine_ctx_var)` resolve to undef (0) because
135+ // the SSA reconstruction can't see backward through the not-yet
136+ // -emitted switch.
137+ builder. switch_to_block ( b_lookup) ;
138+ let exec_start = ctx. exec_start ( builder, machine_ctx_var) ;
139+ let disc = ExecCtxLayout :: load ( builder, ptr_ty, exec_start, ExecCtxLayout :: TEMP_VAL ) ;
140+ let in_range = builder. ins ( ) . icmp_imm ( IntCC :: UnsignedLessThan , disc, n) ;
141+ builder. ins ( ) . brif ( in_range, b_inrange, & [ ] , b_oob, & [ ] ) ;
142+
143+ // In-range: load values[disc], store to TEMP_VAL, continue.
144+ builder. switch_to_block ( b_inrange) ;
145+ let kt_gv = ctx
146+ . module_mut ( )
147+ . declare_data_in_func ( data_id, builder. func ) ;
148+ let kt_base = builder. ins ( ) . global_value ( ptr_ty, kt_gv) ;
149+ let off = builder. ins ( ) . imul_imm ( disc, 8 ) ;
150+ let addr = builder. ins ( ) . iadd ( kt_base, off) ;
151+ let val = builder. ins ( ) . load ( ptr_ty, MemFlags :: trusted ( ) , addr, 0 ) ;
152+ let exec_start2 = ctx. exec_start ( builder, machine_ctx_var) ;
153+ ExecCtxLayout :: store ( builder, val, exec_start2, ExecCtxLayout :: TEMP_VAL ) ;
154+ let qb = ctx. quantum_block ( ) ;
155+ let after_v = builder. ins ( ) . iconst ( ptr_ty, after_when_id) ;
156+ builder. ins ( ) . jump ( qb, & [ BlockArg :: Value ( after_v) ] ) ;
157+
158+ // OOB: silent fall-through, matches when_no_match_continues semantics.
159+ builder. switch_to_block ( b_oob) ;
160+ let after_v2 = builder. ins ( ) . iconst ( ptr_ty, after_when_id) ;
161+ builder. ins ( ) . jump ( qb, & [ BlockArg :: Value ( after_v2) ] ) ;
162+
163+ outer_switch. set_entry ( disc_done_id as u128 , b_lookup) ;
164+ Ok ( after_when_id)
165+ }
166+
20167enum WhenBranchType {
21168 Simple ,
22169 Ptr ,
@@ -59,6 +206,36 @@ pub fn compile<'a>(
59206 positions. into_iter ( ) . next ( )
60207 } ;
61208
209+ // #079 — Detect the K-table fast path BEFORE creating the
210+ // N-way-switch helper blocks: those blocks would otherwise be
211+ // dangling (created but never filled) when we return early.
212+ if let Some ( values) = try_build_k_table ( & branches) {
213+ let disc_done_id = match compile_expr (
214+ ctx,
215+ builder,
216+ machine_ctx_var,
217+ block_id,
218+ outer_switch,
219+ state,
220+ cond_expr,
221+ None ,
222+ None ,
223+ false ,
224+ ) ? {
225+ StmtOutcome :: Continue ( id) => id,
226+ other => bail ! ( "`when` discriminant cannot be a terminal: {:?}" , other) ,
227+ } ;
228+ let after_when_id = emit_k_table (
229+ ctx,
230+ builder,
231+ machine_ctx_var,
232+ outer_switch,
233+ disc_done_id,
234+ & values,
235+ ) ?;
236+ return Ok ( StmtOutcome :: Continue ( after_when_id) ) ;
237+ }
238+
62239 let b_check = builder. create_block ( ) ;
63240 let b_ret: Vec < _ > = ( 0 ..branches. len ( ) )
64241 . map ( |_| builder. create_block ( ) )
0 commit comments