diff --git a/examples/sensor.mbas b/examples/sensor.mbas index 59e6bd1..a45c832 100644 --- a/examples/sensor.mbas +++ b/examples/sensor.mbas @@ -1 +1,2 @@ health = SENSOR(@unit, @health) +health2 = @unit.health diff --git a/src/config/builtin_function.rs b/src/config/builtin_function.rs index a672755..7845df3 100644 --- a/src/config/builtin_function.rs +++ b/src/config/builtin_function.rs @@ -1,4 +1,4 @@ -use crate::{parse::ParseError, parse::ParseErrorKind, prelude::*, translate::Namer}; +use crate::{parse::ParseError, parse::ParseErrorKind, prelude::*, translate::{Namer, translate_expression}}; pub trait BuiltinFunction { fn validate_args( @@ -9,16 +9,13 @@ pub trait BuiltinFunction { fn translate( &self, - arg_names: &[String], + arguments: &[BasicAstExpression], namer: &mut Namer, config: &Config, target_name: String, ) -> MindustryProgram; } -#[derive(Clone, Copy)] -pub struct Sensor; - macro_rules! expect_n_args { ( $args:expr, $call_span:expr, $expected:expr, $name:expr ) => { if $args.len() != $expected { @@ -30,6 +27,21 @@ macro_rules! expect_n_args { }; } +fn translate_arguments(arguments: &[BasicAstExpression], namer: &mut Namer, config: &Config) -> (Vec, MindustryProgram) { + let names = (0..arguments.len()).map(|_| namer.temporary()).collect::>(); + + let mut res = MindustryProgram::new(); + + for (index, arg) in arguments.iter().enumerate() { + res.append(&mut translate_expression(arg, namer, names[index].clone(), config)); + } + + (names, res) +} + +#[derive(Clone, Copy)] +pub struct Sensor; + impl BuiltinFunction for Sensor { fn validate_args( &self, @@ -42,16 +54,67 @@ impl BuiltinFunction for Sensor { fn translate( &self, - arg_names: &[String], - _namer: &mut Namer, - _config: &Config, + arguments: &[BasicAstExpression], + namer: &mut Namer, + config: &Config, target_name: String, ) -> MindustryProgram { - vec![MindustryOperation::Sensor { + let (arg_names, mut res) = translate_arguments(arguments, namer, config); + + res.push(MindustryOperation::Sensor { out_name: target_name, object: Operand::Variable(arg_names[0].clone()), key: Operand::Variable(arg_names[1].clone()), - }] - .into() + }); + + res + } +} + +/// An alternative to `Sensor` that uses a key for the second argument, instead of a variable that is dereferenced at runtime. +/// Used for the dot (`.`) operator, which allows users to write `@unit.health` instead of `@unit.@health` +#[derive(Clone, Copy)] +pub struct SensorOperator; + +impl BuiltinFunction for SensorOperator { + fn validate_args( + &self, + args: &[BasicAstExpression], + call_span: Position, + ) -> Result<(), ParseError> { + // Should only raise an error if a user calls the function manually + expect_n_args!(args, call_span, 2, "__SENSOR_OPERATOR"); + + match &args[1] { + BasicAstExpression::Variable(_name) => Ok(()), + other => { + Err(ParseError::new( + ParseErrorKind::InvalidArgument(other.clone()), + call_span + )) + } + } + } + + fn translate( + &self, + arguments: &[BasicAstExpression], + namer: &mut Namer, + config: &Config, + target_name: String, + ) -> MindustryProgram { + let object_name = namer.temporary(); + let mut res = translate_expression(&arguments[0], namer, object_name.clone(), config); + let BasicAstExpression::Variable(key) = &arguments[1] else { + unreachable!("Expected second argument to __sensor_operator to be a variable"); + }; + + res.push(MindustryOperation::Sensor { + out_name: target_name, + object: Operand::Variable(object_name), + key: Operand::Variable(format!("@{}", key)) + }); + + res } } diff --git a/src/config/mod.rs b/src/config/mod.rs index 158bd09..e19cbb8 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -127,6 +127,7 @@ impl Default for Config { let mut builtin_functions: HashMap> = HashMap::new(); builtin_functions.insert(String::from("sensor"), Box::new(Sensor)); + builtin_functions.insert(String::from("__sensor_operator"), Box::new(SensorOperator)); Self { builtin_routines: HashMap::from([ diff --git a/src/parse/ast.rs b/src/parse/ast.rs index 8e126a0..8cf76fb 100644 --- a/src/parse/ast.rs +++ b/src/parse/ast.rs @@ -503,7 +503,7 @@ pub(crate) fn parse_expression( )) } } else if let Some(binary_operator) = - BasicOperator::from_fn_name(fn_name_lowercase.as_str()) + Operator::from_fn_name(fn_name_lowercase.as_str()) { if arguments.len() != 2 { Err(ParseError::new( @@ -580,7 +580,7 @@ pub(crate) fn parse_expression( if operator.precedence() < min_precedence { break; } - tokens.take(1); + let operator_pos = tokens.take(1)[0].1; let mut rhs = parse_expression_item(tokens, config)?; while let Some(&BasicToken::Operator(sub_operator)) = peek(tokens) { if sub_operator.precedence() > operator.precedence() { @@ -590,7 +590,29 @@ pub(crate) fn parse_expression( } } - ast = BasicAstExpression::Binary(operator, Box::new(ast), Box::new(rhs)); + ast = match operator { + BasicOperator::Operator(op) => { + BasicAstExpression::Binary(op, Box::new(ast), Box::new(rhs)) + } + BasicOperator::Sensor => { + if let Some(fn_config) = config.builtin_functions.get("__sensor_operator") { + let arguments = vec![ + ast, + rhs + ]; + fn_config.validate_args(&arguments, operator_pos)?; + + BasicAstExpression::BuiltinFunction(String::from("__sensor_operator"), arguments) + } else { + return Err(ParseError::new( + ParseErrorKind::UnexpectedToken(BasicToken::Operator(BasicOperator::Sensor)), + operator_pos + )); + } + } + }; + + } Ok(ast) diff --git a/src/repr/basic.rs b/src/repr/basic.rs index 63814ce..fed34e5 100644 --- a/src/repr/basic.rs +++ b/src/repr/basic.rs @@ -7,7 +7,7 @@ pub enum BasicAstExpression { Variable(String), String(String), Binary( - BasicOperator, + Operator, Box, Box, ), diff --git a/src/repr/operator.rs b/src/repr/operator.rs index 60abe17..7da5aed 100644 --- a/src/repr/operator.rs +++ b/src/repr/operator.rs @@ -35,15 +35,6 @@ impl BasicOperator { Self::Operator(op) => op.precedence(), } } - - pub(crate) fn from_fn_name(raw: &str) -> Option { - match raw { - "max" => Some(Self::Operator(Operator::Max)), - "min" => Some(Self::Operator(Operator::Min)), - "pow" => Some(Self::Operator(Operator::Pow)), - _ => None, - } - } } impl Operator { @@ -58,6 +49,15 @@ impl Operator { _ => 128, } } + + pub(crate) fn from_fn_name(raw: &str) -> Option { + match raw { + "max" => Some(Self::Max), + "min" => Some(Self::Min), + "pow" => Some(Self::Pow), + _ => None, + } + } } impl TryFrom for Operator { diff --git a/src/translate/mod.rs b/src/translate/mod.rs index 88e45df..cb5b81d 100644 --- a/src/translate/mod.rs +++ b/src/translate/mod.rs @@ -23,20 +23,20 @@ impl Default for Namer { } impl Namer { - fn temporary(&mut self) -> String { + pub fn temporary(&mut self) -> String { let res = format!("{}__tmp_{}", self.prefix, self.var_index); self.var_index += 1; res } - fn label(&mut self, suffix: &str) -> String { + pub fn label(&mut self, suffix: &str) -> String { let res = format!("{}__label_{}_{}", self.prefix, self.label_index, suffix); self.label_index += 1; res } } -fn translate_expression( +pub(crate) fn translate_expression( expression: &BasicAstExpression, namer: &mut Namer, target_name: String, @@ -68,23 +68,12 @@ fn translate_expression( res.append(&mut right); - match op { - BasicOperator::Operator(op) => { - res.push(MindustryOperation::Operator( - target_name, - *op, - Operand::Variable(left_name), - Operand::Variable(right_name), - )); - } - BasicOperator::Sensor => { - res.push(MindustryOperation::Sensor { - out_name: target_name, - object: Operand::Variable(left_name), - key: Operand::Variable(right_name), - }); - } - } + res.push(MindustryOperation::Operator( + target_name, + *op, + Operand::Variable(left_name), + Operand::Variable(right_name), + )); res } @@ -101,28 +90,11 @@ fn translate_expression( res } BasicAstExpression::BuiltinFunction(name, arguments) => { - let names = (0..arguments.len()) - .map(|_| namer.temporary()) - .collect::>(); - - let mut res = MindustryProgram::new(); - - for (index, arg) in arguments.iter().enumerate() { - res.append(&mut translate_expression( - arg, - namer, - names[index].clone(), - config, - )); - } - let Some(fn_config) = config.builtin_functions.get(name) else { unreachable!("Builtin function {} not found", name); }; - res.append(&mut fn_config.translate(&names, namer, config, target_name)); - - res + fn_config.translate(arguments, namer, config, target_name) } } } diff --git a/tests/examples/sensor.0.mlog b/tests/examples/sensor.0.mlog index 99c1669..a290613 100644 --- a/tests/examples/sensor.0.mlog +++ b/tests/examples/sensor.0.mlog @@ -1,3 +1,5 @@ set main__tmp_0 @unit set main__tmp_1 @health sensor health main__tmp_0 main__tmp_1 +set main__tmp_2 @unit +sensor health2 main__tmp_2 @health diff --git a/tests/examples/sensor.1.mlog b/tests/examples/sensor.1.mlog index a0c992d..93654a7 100644 --- a/tests/examples/sensor.1.mlog +++ b/tests/examples/sensor.1.mlog @@ -1,2 +1,4 @@ set main__tmp_0 @unit sensor health main__tmp_0 @health +set main__tmp_2 @unit +sensor health2 main__tmp_2 @health