use crate::parser::{DuckyScript, Statements, Expr, Value, Operator, Lit, Variable, Function, BuiltinFn};
use std::{sync::{mpsc::{self, TryRecvError}, atomic::{AtomicBool, Ordering}, Arc}, collections::{VecDeque, BTreeMap}, thread};
use crate::parser::Statement;
pub trait VmHardwareRunner:Sized {
fn led_on(&mut self){}
fn led_off(&mut self){}
fn led_r(&mut self){}
fn led_g(&mut self){}
fn delay(&mut self, _: u32){}
fn script_finished(&mut self){}
}
#[derive(Debug, Clone)]
struct VmInstructionStack {
statements_stack: Vec<VecDeque<Statement>>,
statements: VecDeque<Statement>,
}
impl VmInstructionStack{
fn new(statements: Statements ) -> Self {
Self {
statements: VecDeque::from_iter(statements),
statements_stack: Vec::new()
}
}
pub fn peek_instruction(&mut self) -> Option<&Statement>{
self.statements.front()
}
pub fn next_instruction(&mut self) -> Option<Statement> {
loop {
if let Some(statement) = self.statements.pop_front() {
return Some(statement)
}else{
if let Some(s) = self.statements_stack.pop() {
self.statements = s;
}else{
return None
}
}
}
}
pub fn push_instructions(&mut self, statements: Statements) {
let mut ins = VecDeque::from_iter(statements);
std::mem::swap(&mut self.statements, &mut ins);
self.statements_stack.push(ins);
}
pub fn clear_instructions(&mut self) {
self.statements.clear();
self.statements_stack.clear();
}
}
#[derive(Default, Debug, Clone)]
pub struct VmState{
waiting_for_button_state: Option<Arc<AtomicBool>>,
variables: BTreeMap<String, Lit>,
functions: BTreeMap<String, Statements>,
button_enabled: bool
}
#[derive(Debug, Clone)]
pub struct DuckyScriptVm {
ins: VmInstructionStack,
state: VmState,
button_def: Option<Statements>
}
impl DuckyScriptVm {
pub fn new(script: DuckyScript) -> Self {
Self {
ins: VmInstructionStack::new(script.0) ,
state: VmState{
button_enabled: true,
..Default::default()
},
button_def: None
}
}
pub fn is_waiting_for_button(&self) -> bool {
self.state.waiting_for_button_state.is_some()
}
pub fn press_button(&mut self) {
if !self.state.button_enabled {
return;
}
if let Some(b) = &self.state.waiting_for_button_state {
b.store(true, Ordering::Relaxed);
}else if let Some(button_fn) = &self.button_def {
self.ins.push_instructions(button_fn.clone());
}
}
}
impl DuckyScriptVm {
pub fn run_one(&mut self, runner: &mut impl VmHardwareRunner) -> bool {
if let Some(b) = &self.state.waiting_for_button_state {
if b.load(Ordering::Relaxed) {
return true;
}
}
let while_block = if let Some(Statement::WhileBlock(cond, st)) = self.ins.peek_instruction() {
Some((cond.clone(), st.clone()))
}else{
None
};
if let Some((cond, st)) = while_block {
if let Lit::Bool(true) = self.expr_to_lit(&cond) {
self.ins.push_instructions(st.clone());
}else{
self.ins.next_instruction();
return true;
}
}
if let Some(statement) = self.ins.next_instruction() {
match statement {
Statement::ButtonDef(s) => {
self.button_def = Some(s)
},
Statement::VariableDef(name, expr) => {
self.state.variables.insert(name, self.expr_to_lit(&expr));
},
Statement::CallFn(f) => {
self.run_function(f, runner);
},
_ => unreachable!()
}
true
}else{
false
}
}
pub fn expr_to_lit(&self, expr: &Expr) -> Lit {
match expr {
Expr::Value(v) => self.value_to_lit(v),
Expr::OperationalExpr(ls, op, rs)
=> self.cond_expr_to_val(ls, op, rs)
}
}
pub fn cond_expr_to_val(&self, ls: &Value, operator: &Operator, rs: &Value) -> Lit {
let ls = self.value_to_lit(ls);
let rs = self.value_to_lit(rs);
match (ls, operator, rs) {
(Lit::Bool(ls), Operator::Equals, Lit::Bool(rs)) => Lit::Bool(ls == rs),
(Lit::Int(ls), Operator::Equals, Lit::Int(rs)) => Lit::Bool(ls == rs),
(Lit::Str(ls), Operator::Equals, Lit::Str(rs)) => Lit::Bool(ls == rs),
(Lit::Int(ls), Operator::GreaterThan, Lit::Int(rs)) => Lit::Bool(ls > rs),
(Lit::Int(ls), Operator::LessThan, Lit::Int(rs)) => Lit::Bool(ls < rs),
(Lit::Int(ls), Operator::Add, Lit::Int(rs)) => Lit::Int(ls + rs),
(Lit::Int(ls), Operator::Subtract, Lit::Int(rs)) => Lit::Int(ls + rs),
(Lit::Str(mut ls), Operator::Add, Lit::Str(rs)) => {
ls.push_str(&rs);
Lit::Str(ls)
},
_ => Lit::Bool(false)
}
}
pub fn value_to_lit(&self, value: &Value) -> Lit {
match value {
Value::Lit(l) => l.clone(),
Value::Variable(v) => self.get_variable_lit(v)
}
}
fn get_variable_lit(&self, v: &Variable) -> Lit {
match v {
Variable::Builtin(b) => {
todo!()
},
Variable::Script(s) => {
if let Some(v) = self.state.variables.get(s) {
v.clone()
}else{
Lit::Bool(false)
}
}
}
}
fn run_function(&mut self, f: Function, runner: &mut impl VmHardwareRunner) {
match f {
Function::UserDefined(f) => {
if let Some(f) = self.state.functions.get(&f) {
self.ins.push_instructions(f.clone());
}
},
Function::Builtin(b) => {
match b {
BuiltinFn::Delay(v) => {
if let Lit::Int(delay_ms) = self.value_to_lit(&v) {
runner.delay(delay_ms);
}
},
BuiltinFn::LedOff => {
runner.led_off()
},
BuiltinFn::LedOn => {
runner.led_on()
},
BuiltinFn::LedG => {
runner.led_g()
},
BuiltinFn::LedR => {
runner.led_r()
},
BuiltinFn::StopPayload => {
self.ins.clear_instructions();
}
}
}
}
}
}
pub enum VmCommand {
ButtonPressed,
}
#[derive(Default, Debug, Clone)]
pub struct DuckyScriptRunner<T: VmHardwareRunner+Send> {
runner: Option<T>,
msg_tx: Option<mpsc::Sender<VmCommand>>
}
impl<T:VmHardwareRunner+Send+'static> DuckyScriptRunner<T> {
pub fn new(runner: T) -> Self {
Self { runner: Some(runner), msg_tx: None }
}
pub fn run(&mut self, script: DuckyScript) {
if let Some(r) = self.runner.take() {
let (tx,rx) = mpsc::channel();
self.msg_tx = Some(tx);
thread::spawn(move ||{
run_vm_thread(script, rx, r);
});
}
}
pub fn send_command(&self, cmd: VmCommand) {
if let Some(c) = &self.msg_tx {
c.send(cmd).ok();
}
}
pub fn press_button(&self) {
self.send_command(VmCommand::ButtonPressed);
}
}
fn run_vm_thread(script: DuckyScript, rx: mpsc::Receiver<VmCommand>, mut runner: impl VmHardwareRunner) {
let mut vm = DuckyScriptVm::new(script);
while vm.run_one(&mut runner) {
let cmd = if vm.is_waiting_for_button(){
match rx.recv() {
Ok(c) => Some(c),
_ => break
}
}else{
match rx.try_recv() {
Ok(c) => Some(c),
Err(TryRecvError::Disconnected) => break,
_ => None
}
};
match cmd {
Some(VmCommand::ButtonPressed) => {
vm.press_button();
},
_ =>{}
}
}
runner.script_finished();
}