Skip to main content

flux_vm/
lib.rs

1use core::result::Result;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4pub enum Fault {
5    StackUnderflow,
6    StackOverflow,
7    GasExhausted,
8    AssertFailed,
9    GuardTrap,
10    CallStackOverflow,
11    CallStackUnderflow,
12    InvalidMemoryAccess,
13}
14
15const STACK_SIZE: usize = 256;
16const MEMORY_SIZE: usize = 65536;
17const CALL_STACK_SIZE: usize = 32;
18
19pub struct FluxVM {
20    stack: [u8; STACK_SIZE],
21    sp: usize,
22    pc: usize,
23    gas: u32,
24    halted: bool,
25    yielded: bool,
26    memory: [u8; MEMORY_SIZE],
27    call_stack: [usize; CALL_STACK_SIZE],
28    csp: usize,
29    guard_reg: u8,
30    last_check_passed: bool,
31}
32
33impl FluxVM {
34    pub fn new(gas: u32) -> Self {
35        Self {
36            stack: [0u8; STACK_SIZE],
37            sp: 0,
38            pc: 0,
39            gas,
40            halted: false,
41            yielded: false,
42            memory: [0u8; MEMORY_SIZE],
43            call_stack: [0usize; CALL_STACK_SIZE],
44            csp: 0,
45            guard_reg: 0,
46            last_check_passed: true,
47        }
48    }
49
50    fn push(&mut self, value: u8) -> Result<(), Fault> {
51        if self.sp >= STACK_SIZE { return Err(Fault::StackOverflow); }
52        self.stack[self.sp] = value;
53        self.sp += 1;
54        Ok(())
55    }
56
57    fn pop(&mut self) -> Result<u8, Fault> {
58        if self.sp == 0 { return Err(Fault::StackUnderflow); }
59        self.sp -= 1;
60        Ok(self.stack[self.sp])
61    }
62
63    fn peek(&self) -> Result<u8, Fault> {
64        if self.sp == 0 { return Err(Fault::StackUnderflow); }
65        Ok(self.stack[self.sp - 1])
66    }
67
68    fn read_byte(&self, bytecode: &[u8]) -> Result<u8, Fault> {
69        bytecode.get(self.pc).copied().ok_or(Fault::InvalidMemoryAccess)
70    }
71
72    fn binop<F: FnOnce(u8, u8) -> u8>(&mut self, f: F) -> Result<(), Fault> {
73        let b = self.pop()?;
74        let a = self.pop()?;
75        self.push(f(a, b))
76    }
77
78    fn cmpop<F: FnOnce(u8, u8) -> bool>(&mut self, f: F) -> Result<(), Fault> {
79        let b = self.pop()?;
80        let a = self.pop()?;
81        self.push(if f(a, b) { 1 } else { 0 })
82    }
83
84    pub fn step(&mut self, bytecode: &[u8]) -> Result<bool, Fault> {
85        if self.halted { return Ok(true); }
86        if self.yielded { self.yielded = false; }
87        if self.gas == 0 { return Err(Fault::GasExhausted); }
88        if self.pc >= bytecode.len() { return Ok(true); }
89
90        let op = bytecode[self.pc];
91        self.pc += 1;
92        self.gas -= 1;
93
94        match op {
95            // === Stack Operations ===
96            0x00 => { // PUSH val
97                let v = self.read_byte(bytecode)?;
98                self.pc += 1;
99                self.push(v)?;
100            }
101            0x01 => { self.pop()?; } // POP
102            0x02 => { // DUP
103                let v = self.peek()?;
104                self.push(v)?;
105            }
106            0x03 => { // SWAP
107                let b = self.pop()?;
108                let a = self.pop()?;
109                self.push(b)?;
110                self.push(a)?;
111            }
112
113            // === Memory ===
114            0x04 => { // LOAD addr
115                let addr = self.read_byte(bytecode)? as usize;
116                self.pc += 1;
117                let v = *self.memory.get(addr).ok_or(Fault::InvalidMemoryAccess)?;
118                self.push(v)?;
119            }
120            0x05 => { // STORE addr
121                let addr = self.read_byte(bytecode)? as usize;
122                self.pc += 1;
123                let v = self.pop()?;
124                *self.memory.get_mut(addr).ok_or(Fault::InvalidMemoryAccess)? = v;
125            }
126
127            // === Arithmetic ===
128            0x06 => self.binop(|a, b| a.wrapping_add(b))?, // ADD
129            0x07 => self.binop(|a, b| a.wrapping_sub(b))?, // SUB
130            0x08 => self.binop(|a, b| a.wrapping_mul(b))?, // MUL
131
132            // === Bitwise ===
133            0x09 => self.binop(|a, b| a & b)?, // AND
134            0x0A => self.binop(|a, b| a | b)?, // OR
135            0x0B => self.binop(|a, b| a ^ b)?, // XOR
136            0x0C => { let a = self.pop()?; self.push(!a)?; } // NOT
137            0x0D => { let a = self.pop()?; self.push(a << 1)?; } // SHL
138            0x0E => { let a = self.pop()?; self.push(a >> 1)?; } // SHR
139
140            // === Comparison ===
141            0x0F => self.cmpop(|a, b| a == b)?, // EQ
142            0x10 => self.cmpop(|a, b| a != b)?, // NEQ
143            0x11 => self.cmpop(|a, b| a < b)?,  // LT
144            0x12 => self.cmpop(|a, b| a > b)?,  // GT
145            0x13 => self.cmpop(|a, b| a <= b)?, // LTE
146            0x14 => self.cmpop(|a, b| a >= b)?, // GTE
147
148            // === Control Flow ===
149            0x15 => { // JUMP addr
150                let addr = self.read_byte(bytecode)? as usize;
151                self.pc = addr;
152            }
153            0x16 => { // JZ addr
154                let addr = self.read_byte(bytecode)? as usize;
155                self.pc += 1; // consume operand even if we jump
156                let v = self.pop()?;
157                if v == 0 { self.pc = addr; }
158            }
159            0x17 => { // JNZ addr
160                let addr = self.read_byte(bytecode)? as usize;
161                self.pc += 1;
162                let v = self.pop()?;
163                if v != 0 { self.pc = addr; }
164            }
165            0x18 => { // CALL addr
166                let addr = self.read_byte(bytecode)? as usize;
167                self.pc += 1;
168                if self.csp >= CALL_STACK_SIZE { return Err(Fault::CallStackOverflow); }
169                self.call_stack[self.csp] = self.pc;
170                self.csp += 1;
171                self.pc = addr;
172            }
173            0x19 => { // RET
174                if self.csp == 0 { return Err(Fault::CallStackUnderflow); }
175                self.csp -= 1;
176                self.pc = self.call_stack[self.csp];
177            }
178
179            // === Execution Control ===
180            0x1A => { self.halted = true; } // HALT
181            0x1B => { // ASSERT
182                let v = self.pop()?;
183                self.last_check_passed = v != 0;
184                if v == 0 { return Err(Fault::AssertFailed); }
185            }
186
187            // === Constraint Checking ===
188            0x1C => { // CHECK_DOMAIN mask
189                let mask = self.read_byte(bytecode)?;
190                self.pc += 1;
191                let v = self.pop()?;
192                let result = v & mask;
193                self.last_check_passed = result != 0;
194                self.push(result)?;
195            }
196            0x1D => { // BITMASK_RANGE lo hi
197                let lo = self.read_byte(bytecode)?;
198                self.pc += 1;
199                let hi = self.read_byte(bytecode)?;
200                self.pc += 1;
201                let v = self.pop()?;
202                let in_range = v >= lo && v <= hi;
203                self.last_check_passed = in_range;
204                self.push(if in_range { 1 } else { 0 })?;
205            }
206            0x1E => { // LOAD_GUARD
207                self.push(self.guard_reg)?;
208            }
209            0x1F => { // MERKLE_VERIFY (simplified: pop 4 bytes, compare to stored)
210                let _b3 = self.pop()?;
211                let _b2 = self.pop()?;
212                let _b1 = self.pop()?;
213                let _b0 = self.pop()?;
214                // Simplified: always pass for now
215                self.last_check_passed = true;
216                self.push(1)?;
217            }
218            0x20 => { return Err(Fault::GuardTrap); } // GUARD_TRAP
219
220            // === Hash/Crypto ===
221            0x21 => { // CRC32 (simplified: XOR-fold stack)
222                let mut acc: u8 = 0;
223                for i in 0..self.sp {
224                    acc ^= self.stack[i];
225                }
226                self.push(acc)?;
227            }
228            0x22 => { // PUSH_HASH hi lo
229                let hi = self.read_byte(bytecode)?;
230                let lo = self.read_byte(bytecode)?;
231                self.pc += 2;
232                self.push(hi)?;
233                self.push(lo)?;
234            }
235            0x23 => { // XNOR_POPCOUNT
236                let b = self.pop()?;
237                let a = self.pop()?;
238                let xnor = !(a ^ b);
239                let count = xnor.count_ones() as u8;
240                self.push(count)?;
241            }
242
243            // === Extended Comparison ===
244            0x24 => self.cmpop(|a, b| a >= b)?, // CMP_GE (same as GTE)
245            0x25 => { // CARRY_LT: pop a,b, push 1 if a < b (unsigned)
246                let b = self.pop()?;
247                let a = self.pop()?;
248                self.push(if a < b { 1 } else { 0 })?;
249            }
250            0x26 => { // JFAIL addr
251                let addr = self.read_byte(bytecode)? as usize;
252                self.pc += 1;
253                if !self.last_check_passed { self.pc = addr; }
254            }
255
256            // === Misc ===
257            0x27 => {} // NOP
258            0x28 => { self.sp = 0; } // FLUSH
259            0x29 => { self.yielded = true; } // YIELD
260
261            _ => {} // Unknown opcodes are NOP
262        }
263
264        Ok(self.halted)
265    }
266
267    pub fn execute(&mut self, bytecode: &[u8], max_steps: usize) -> Result<(), Vec<Fault>> {
268        for _ in 0..max_steps {
269            match self.step(bytecode) {
270                Ok(true) => return Ok(()),
271                Ok(false) => continue,
272                Err(f) => return Err(vec![f]),
273            }
274        }
275        Ok(())
276    }
277
278    // === Public Accessors ===
279    pub fn is_halted(&self) -> bool { self.halted }
280    pub fn is_yielded(&self) -> bool { self.yielded }
281    pub fn stack_top(&self) -> Option<u8> {
282        if self.sp > 0 { Some(self.stack[self.sp - 1]) } else { None }
283    }
284    pub fn stack_len(&self) -> usize { self.sp }
285    pub fn gas_remaining(&self) -> u32 { self.gas }
286    pub fn pc(&self) -> usize { self.pc }
287    pub fn get_memory(&self, addr: usize) -> Option<u8> { self.memory.get(addr).copied() }
288    pub fn set_memory(&mut self, addr: usize, val: u8) { if addr < MEMORY_SIZE { self.memory[addr] = val; } }
289    pub fn set_guard(&mut self, val: u8) { self.guard_reg = val; }
290    pub fn last_check_passed(&self) -> bool { self.last_check_passed }
291}
292
293#[cfg(test)]
294mod tests {
295    use super::*;
296
297    #[test]
298    fn test_push_add_halt() {
299        let mut vm = FluxVM::new(100);
300        vm.execute(&[0x00, 3, 0x00, 4, 0x06, 0x1A], 100).unwrap();
301        assert!(vm.halted);
302        assert_eq!(vm.stack[0], 7);
303    }
304
305    #[test]
306    fn test_assert_fail() {
307        let mut vm = FluxVM::new(100);
308        let result = vm.execute(&[0x00, 0, 0x1B], 100);
309        assert!(result.is_err());
310    }
311
312    #[test]
313    fn test_guard_trap() {
314        let mut vm = FluxVM::new(100);
315        let result = vm.execute(&[0x20], 100);
316        assert!(matches!(result, Err(f) if f[0] == Fault::GuardTrap));
317    }
318
319    #[test]
320    fn test_gas_exhaustion() {
321        let mut vm = FluxVM::new(2);
322        let result = vm.execute(&[0x00, 1, 0x00, 2, 0x00, 3], 100);
323        assert!(matches!(result, Err(f) if f[0] == Fault::GasExhausted));
324    }
325
326    #[test]
327    fn test_jump_control_flow() {
328        let mut vm = FluxVM::new(100);
329        vm.execute(&[0x00, 0, 0x16, 7, 0x00, 99, 0x1A, 0x00, 42, 0x1A], 100).unwrap();
330        assert_eq!(vm.stack[0], 42);
331    }
332
333    // === New opcode tests ===
334
335    #[test]
336    fn test_dup() {
337        let mut vm = FluxVM::new(100);
338        // PUSH 7, DUP, HALT → stack has [7, 7]
339        vm.execute(&[0x00, 7, 0x02, 0x1A], 100).unwrap();
340        assert_eq!(vm.stack_len(), 2);
341        assert_eq!(vm.stack_top(), Some(7));
342    }
343
344    #[test]
345    fn test_swap() {
346        let mut vm = FluxVM::new(100);
347        // PUSH 3, PUSH 5, SWAP → [5, 3]
348        vm.execute(&[0x00, 3, 0x00, 5, 0x03, 0x1A], 100).unwrap();
349        assert_eq!(vm.stack_top(), Some(3)); // top is 3 (was bottom)
350    }
351
352    #[test]
353    fn test_memory_load_store() {
354        let mut vm = FluxVM::new(100);
355        // PUSH 42, STORE 100, LOAD 100, HALT
356        vm.execute(&[0x00, 42, 0x05, 100, 0x04, 100, 0x1A], 100).unwrap();
357        assert_eq!(vm.stack_top(), Some(42));
358    }
359
360    #[test]
361    fn test_shl_shr() {
362        let mut vm = FluxVM::new(100);
363        // PUSH 1, SHL → 2, SHL → 4, SHR → 2
364        vm.execute(&[0x00, 1, 0x0D, 0x0D, 0x0E, 0x1A], 100).unwrap();
365        assert_eq!(vm.stack_top(), Some(2));
366    }
367
368    #[test]
369    fn test_lte_gte() {
370        let mut vm = FluxVM::new(100);
371        // 5 <= 5 → 1
372        vm.execute(&[0x00, 5, 0x00, 5, 0x13, 0x1A], 100).unwrap();
373        assert_eq!(vm.stack_top(), Some(1));
374    }
375
376    #[test]
377    fn test_jnz() {
378        let mut vm = FluxVM::new(100);
379        // PUSH 1, JNZ 7, PUSH 99, HALT, PUSH 42, HALT
380        vm.execute(&[0x00, 1, 0x17, 7, 0x00, 99, 0x1A, 0x00, 42, 0x1A], 100).unwrap();
381        assert_eq!(vm.stack_top(), Some(42)); // took the jump
382    }
383
384    #[test]
385    fn test_call_ret() {
386        let mut vm = FluxVM::new(100);
387        // CALL 5, HALT, PUSH 42, RET
388        vm.execute(&[0x18, 5, 0x1A, 0x00, 42, 0x19, 0x1A], 100).unwrap();
389        // After CALL 5, runs PUSH 42, RET, continues to... 
390        // Actually CALL pushes return addr (2), jumps to 5
391        // addr 5: RET → returns to addr 2 (the HALT)
392        assert!(vm.is_halted());
393    }
394
395    #[test]
396    fn test_check_domain() {
397        let mut vm = FluxVM::new(100);
398        // PUSH 0x42, CHECK_DOMAIN 0x0F → 0x42 & 0x0F = 0x02
399        vm.execute(&[0x00, 0x42, 0x1C, 0x0F, 0x1A], 100).unwrap();
400        assert_eq!(vm.stack_top(), Some(0x02));
401        assert!(vm.last_check_passed());
402    }
403
404    #[test]
405    fn test_bitmask_range() {
406        let mut vm = FluxVM::new(100);
407        // PUSH 50, BITMASK_RANGE 0 100 → in range, push 1
408        vm.execute(&[0x00, 50, 0x1D, 0, 100, 0x1A], 100).unwrap();
409        assert_eq!(vm.stack_top(), Some(1));
410        assert!(vm.last_check_passed());
411    }
412
413    #[test]
414    fn test_bitmask_range_fail() {
415        let mut vm = FluxVM::new(100);
416        // PUSH 200, BITMASK_RANGE 0 100 → out of range, push 0
417        vm.execute(&[0x00, 200, 0x1D, 0, 100, 0x1A], 100).unwrap();
418        assert_eq!(vm.stack_top(), Some(0));
419        assert!(!vm.last_check_passed());
420    }
421
422    #[test]
423    fn test_xnor_popcount() {
424        let mut vm = FluxVM::new(100);
425        // PUSH 0xFF, PUSH 0xFF, XNOR_POPCOUNT → 8 bits match
426        vm.execute(&[0x00, 0xFF, 0x00, 0xFF, 0x23, 0x1A], 100).unwrap();
427        assert_eq!(vm.stack_top(), Some(8));
428    }
429
430    #[test]
431    fn test_carry_lt() {
432        let mut vm = FluxVM::new(100);
433        // PUSH 3, PUSH 5, CARRY_LT → 3 < 5 = true → 1
434        vm.execute(&[0x00, 3, 0x00, 5, 0x25, 0x1A], 100).unwrap();
435        assert_eq!(vm.stack_top(), Some(1));
436    }
437
438    #[test]
439    fn test_jfail() {
440        let mut vm = FluxVM::new(100);
441        // PUSH 50, BITMASK_RANGE 0 100 → pass, JFAIL should NOT jump
442        vm.execute(&[0x00, 50, 0x1D, 0, 100, 0x26, 10, 0x00, 77, 0x1A, 0x20], 100).unwrap();
443        assert_eq!(vm.stack_top(), Some(77)); // didn't jump to GUARD_TRAP
444    }
445
446    #[test]
447    fn test_flush() {
448        let mut vm = FluxVM::new(100);
449        // PUSH 1, PUSH 2, PUSH 3, FLUSH, PUSH 42, HALT
450        vm.execute(&[0x00, 1, 0x00, 2, 0x00, 3, 0x28, 0x00, 42, 0x1A], 100).unwrap();
451        assert_eq!(vm.stack_len(), 1);
452        assert_eq!(vm.stack_top(), Some(42));
453    }
454
455    #[test]
456    fn test_yield() {
457        let mut vm = FluxVM::new(100);
458        vm.step(&[0x29]).unwrap(); // YIELD
459        assert!(vm.is_yielded());
460        assert!(!vm.is_halted());
461    }
462
463    #[test]
464    fn test_load_guard() {
465        let mut vm = FluxVM::new(100);
466        vm.set_guard(0xAB);
467        vm.execute(&[0x1E, 0x1A], 100).unwrap(); // LOAD_GUARD, HALT
468        assert_eq!(vm.stack_top(), Some(0xAB));
469    }
470
471    // === Certification test vectors (15 programs) ===
472
473    #[test]
474    fn cert_identity() {
475        let mut vm = FluxVM::new(100);
476        vm.execute(&[0x00, 42, 0x1A], 100).unwrap();
477        assert_eq!(vm.stack[0], 42);
478    }
479
480    #[test]
481    fn cert_add() {
482        let mut vm = FluxVM::new(100);
483        vm.execute(&[0x00, 3, 0x00, 4, 0x06, 0x1A], 100).unwrap();
484        assert_eq!(vm.stack[0], 7);
485    }
486
487    #[test]
488    fn cert_mul() {
489        let mut vm = FluxVM::new(100);
490        vm.execute(&[0x00, 6, 0x00, 7, 0x08, 0x1A], 100).unwrap();
491        assert_eq!(vm.stack[0], 42);
492    }
493
494    #[test]
495    fn cert_sub() {
496        let mut vm = FluxVM::new(100);
497        vm.execute(&[0x00, 10, 0x00, 3, 0x07, 0x1A], 100).unwrap();
498        assert_eq!(vm.stack[0], 7);
499    }
500
501    #[test]
502    fn cert_and_mask() {
503        let mut vm = FluxVM::new(100);
504        vm.execute(&[0x00, 0xFF, 0x00, 0x0F, 0x09, 0x1A], 100).unwrap();
505        assert_eq!(vm.stack[0], 0x0F);
506    }
507
508    #[test]
509    fn cert_or() {
510        let mut vm = FluxVM::new(100);
511        vm.execute(&[0x00, 0xF0, 0x00, 0x0F, 0x0A, 0x1A], 100).unwrap();
512        assert_eq!(vm.stack[0], 0xFF);
513    }
514
515    #[test]
516    fn cert_xor() {
517        let mut vm = FluxVM::new(100);
518        vm.execute(&[0x00, 0xAA, 0x00, 0x55, 0x0B, 0x1A], 100).unwrap();
519        assert_eq!(vm.stack[0], 0xFF);
520    }
521
522    #[test]
523    fn cert_not() {
524        let mut vm = FluxVM::new(100);
525        vm.execute(&[0x00, 0x00, 0x0C, 0x1A], 100).unwrap();
526        assert_eq!(vm.stack[0], 0xFF);
527    }
528
529    #[test]
530    fn cert_eq() {
531        let mut vm = FluxVM::new(100);
532        vm.execute(&[0x00, 7, 0x00, 7, 0x0F, 0x1A], 100).unwrap();
533        assert_eq!(vm.stack[0], 1);
534    }
535
536    #[test]
537    fn cert_neq() {
538        let mut vm = FluxVM::new(100);
539        vm.execute(&[0x00, 3, 0x00, 5, 0x10, 0x1A], 100).unwrap();
540        assert_eq!(vm.stack[0], 1);
541    }
542
543    #[test]
544    fn cert_lt() {
545        let mut vm = FluxVM::new(100);
546        vm.execute(&[0x00, 3, 0x00, 5, 0x11, 0x1A], 100).unwrap();
547        assert_eq!(vm.stack[0], 1);
548    }
549
550    #[test]
551    fn cert_gt() {
552        let mut vm = FluxVM::new(100);
553        vm.execute(&[0x00, 5, 0x00, 3, 0x12, 0x1A], 100).unwrap();
554        assert_eq!(vm.stack[0], 1);
555    }
556
557    #[test]
558    fn cert_jz_skip() {
559        let mut vm = FluxVM::new(100);
560        vm.execute(&[0x00, 0, 0x16, 7, 0x00, 99, 0x1A, 0x00, 42, 0x1A], 100).unwrap();
561        assert_eq!(vm.stack[0], 42);
562    }
563
564    #[test]
565    fn cert_assert_pass() {
566        let mut vm = FluxVM::new(100);
567        vm.execute(&[0x00, 1, 0x1B, 0x00, 77, 0x1A], 100).unwrap();
568        assert_eq!(vm.stack[0], 77);
569    }
570
571    #[test]
572    fn cert_nops() {
573        let mut vm = FluxVM::new(100);
574        vm.execute(&[0x27, 0x27, 0x00, 13, 0x1A], 100).unwrap();
575        assert_eq!(vm.stack[0], 13);
576    }
577
578    // === Extended tests ===
579
580    #[test]
581    fn test_bitwise_and_mask() {
582        let mut vm = FluxVM::new(100);
583        vm.execute(&[0x00, 0xFF, 0x00, 0x0F, 0x09, 0x1A], 100).unwrap();
584        assert_eq!(vm.stack[0], 0x0F);
585    }
586
587    #[test]
588    fn test_domain_check_pass() {
589        let mut vm = FluxVM::new(100);
590        vm.execute(&[0x00, 0x42, 0x00, 0x0F, 0x09, 0x00, 0, 0x0F, 0x0C, 0x1B, 0x1A], 100).unwrap();
591        assert!(vm.is_halted());
592    }
593
594    #[test]
595    fn test_xor_swap() {
596        let mut vm = FluxVM::new(100);
597        vm.execute(&[0x00, 0xAA, 0x00, 0x55, 0x0B, 0x1A], 100).unwrap();
598        assert_eq!(vm.stack[0], 0xFF);
599    }
600
601    #[test]
602    fn test_comparison_gt() {
603        let mut vm = FluxVM::new(100);
604        vm.execute(&[0x00, 10, 0x00, 5, 0x12, 0x1B, 0x1A], 100).unwrap();
605        assert!(vm.is_halted());
606    }
607
608    #[test]
609    fn test_nested_if_else() {
610        let mut vm = FluxVM::new(100);
611        vm.execute(&[0x00, 5, 0x00, 3, 0x12, 0x16, 10, 0x00, 42, 0x1A, 0x00, 99, 0x1A], 100).unwrap();
612        assert_eq!(vm.stack[0], 42);
613    }
614
615    #[test]
616    fn test_sub_and_assert() {
617        let mut vm = FluxVM::new(100);
618        vm.execute(&[0x00, 10, 0x00, 3, 0x07, 0x00, 7, 0x0F, 0x1B, 0x1A], 100).unwrap();
619        assert!(vm.is_halted());
620    }
621}