✨ Start work on builtin functions

examples/difficulty.mbas

@@ -19,19 +19,19 @@ PRINT_FLUSH(message1)
 
 difficulty = 0
 
-easy = switch1.@enabled
+easy = SENSOR(switch1, @enabled)
 IF easy THEN
     difficulty = 1
     CONTROL(enabled, switch1, false)
 END IF
 
-medium = switch2.@enabled
+medium = SENSOR(switch2, @enabled)
 IF medium THEN
     difficulty = 2
     CONTROL(enabled, switch2, false)
 END IF
 
-hard = switch3.@enabled
+hard = SENSOR(switch3, @enabled)
 IF hard THEN
     difficulty = 3
     CONTROL(enabled, switch3, false)

examples/sensor.mbas

@@ -0,0 +1 @@
+health = SENSOR(@unit, @health)

src/config/builtin_function.rs

@@ -0,0 +1,57 @@
+use crate::{parse::ParseError, parse::ParseErrorKind, prelude::*, translate::Namer};
+
+pub trait BuiltinFunction {
+    fn validate_args(
+        &self,
+        args: &[BasicAstExpression],
+        call_span: Position,
+    ) -> Result<(), ParseError>;
+
+    fn translate(
+        &self,
+        arg_names: &[String],
+        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 {
+            return Err(ParseError::new(
+                ParseErrorKind::InvalidArgumentCount(String::from($name), $expected, $args.len()),
+                $call_span,
+            ));
+        }
+    };
+}
+
+impl BuiltinFunction for Sensor {
+    fn validate_args(
+        &self,
+        args: &[BasicAstExpression],
+        call_span: Position,
+    ) -> Result<(), ParseError> {
+        expect_n_args!(args, call_span, 2, "SENSOR");
+        Ok(())
+    }
+
+    fn translate(
+        &self,
+        arg_names: &[String],
+        _namer: &mut Namer,
+        _config: &Config,
+        target_name: String,
+    ) -> MindustryProgram {
+        vec![MindustryOperation::Sensor {
+            out_name: target_name,
+            object: Operand::Variable(arg_names[0].clone()),
+            key: Operand::Variable(arg_names[1].clone()),
+        }]
+        .into()
+    }
+}

src/config/mod.rs

@@ -6,6 +6,9 @@ use crate::{
     repr::position::Position,
 };
 
+mod builtin_function;
+use builtin_function::*;
+
 pub struct Config {
     pub builtin_routines: HashMap<String, (Option<String>, bool, usize)>,
 
@@ -14,6 +17,8 @@ pub struct Config {
         String,
         Box<dyn Fn(Vec<BasicAstExpression>, Position) -> Result<BasicAstInstruction, ParseError>>,
     >,
+
+    pub builtin_functions: HashMap<String, Box<dyn BuiltinFunction>>,
 }
 
 impl Default for Config {
@@ -31,12 +36,12 @@ impl Default for Config {
             };
         }
 
-        let mut special_functions: HashMap<
+        let mut special_routines: HashMap<
             String,
             Box<dyn Fn(Vec<BasicAstExpression>, Position) -> Result<BasicAstInstruction, _>>,
         > = HashMap::new();
 
-        special_functions.insert(
+        special_routines.insert(
             String::from("print_flush_global"),
             Box::new(|arguments: Vec<BasicAstExpression>, position| {
                 let BasicAstExpression::Variable(buffer) = &arguments[0] else {
@@ -77,7 +82,7 @@ impl Default for Config {
             }),
         );
 
-        special_functions.insert(
+        special_routines.insert(
             String::from("control"),
             Box::new(|arguments, position| {
                 let BasicAstExpression::Variable(buffer) = &arguments[0] else {
@@ -119,6 +124,10 @@ impl Default for Config {
             }),
         );
 
+        let mut builtin_functions: HashMap<String, Box<dyn BuiltinFunction>> = HashMap::new();
+
+        builtin_functions.insert(String::from("sensor"), Box::new(Sensor));
+
         Self {
             builtin_routines: HashMap::from([
                 builtin_function!("print_flush", None, false, 1),
@@ -131,7 +140,8 @@ impl Default for Config {
                 // TODO: same thing
                 builtin_function!("spawn", Some("spawn"), false, 6),
             ]),
-            special_routines: special_functions,
+            special_routines,
+            builtin_functions,
         }
     }
 }

src/optimize/constant.rs

@@ -5,11 +5,11 @@ use crate::prelude::*;
 pub fn optimize_constant(program: MindustryProgram) -> MindustryProgram {
     // Returns true if the variable is safe to substitute into an operand
     let is_safe_variable = |name: &str| -> bool {
-        if matches!(name, "@this" | "@thisx" | "@thisy" | "@links") || is_unit_constant(name) {
-            return true;
+        if matches!(name, "@counter" | "@ipt" | "@time" | "@tick" | "@unit") {
+            false
+        } else {
+            true
         }
-
-        !name.starts_with('@')
     };
 
     let res = replace_if(program, |instructions, instruction, use_index| {
@@ -102,7 +102,8 @@ pub fn optimize_constant(program: MindustryProgram) -> MindustryProgram {
     optimize_dead_code(res)
 }
 
-// TODO: add serpulo units
+// TODO: add serpulo units and move to a text file
+#[allow(dead_code)]
 fn is_unit_constant(name: &str) -> bool {
     matches!(
         name,

src/parse/ast.rs

@@ -82,7 +82,7 @@ pub fn build_ast(
                 tokens.take(1);
                 let then_index = find_token_index(&tokens, BasicToken::Then)?;
 
-                let condition = parse_expression(&mut tokens.range(0..then_index))?;
+                let condition = parse_expression(&mut tokens.range(0..then_index), config)?;
 
                 tokens.take(then_index + 1);
 
@@ -130,17 +130,17 @@ pub fn build_ast(
             [BasicToken::For, BasicToken::Name(variable), BasicToken::Assign, ..] => {
                 tokens.take(3);
 
-                let start = parse_expression(&mut tokens)?;
+                let start = parse_expression(&mut tokens, config)?;
 
                 expect_next_token(&mut tokens, &BasicToken::To)?;
                 tokens.take(1);
 
-                let end = parse_expression(&mut tokens)?;
+                let end = parse_expression(&mut tokens, config)?;
 
                 let step = if let Some((BasicToken::Step, _pos)) = tokens.get(0) {
                     tokens.take(1);
 
-                    parse_expression(&mut tokens)?
+                    parse_expression(&mut tokens, config)?
                 } else {
                     BasicAstExpression::Integer(1)
                 };
@@ -180,14 +180,14 @@ pub fn build_ast(
             // == While loops ==
             [BasicToken::While, ..] => {
                 tokens.take(1);
-                let condition = parse_expression(&mut tokens)?;
+                let condition = parse_expression(&mut tokens, config)?;
                 expect_next_token(&tokens, &BasicToken::NewLine)?;
 
                 context_stack.push((Vec::new(), Context::While(condition)));
             }
             [BasicToken::Do, BasicToken::While, ..] => {
                 tokens.take(2);
-                let condition = parse_expression(&mut tokens)?;
+                let condition = parse_expression(&mut tokens, config)?;
                 expect_next_token(&tokens, &BasicToken::NewLine)?;
 
                 context_stack.push((Vec::new(), Context::DoWhile(condition)));
@@ -213,7 +213,7 @@ pub fn build_ast(
             [BasicToken::Loop, BasicToken::While, ..] => {
                 tokens.take(2);
 
-                let condition = parse_expression(&mut tokens)?;
+                let condition = parse_expression(&mut tokens, config)?;
 
                 pop_context!(context_stack, instructions, {
                     Context::Do => {
@@ -271,7 +271,7 @@ pub fn build_ast(
             // == Misc ==
             [BasicToken::Name(variable_name), BasicToken::Assign, ..] => {
                 tokens.take(2);
-                let expression = parse_expression(&mut tokens)?;
+                let expression = parse_expression(&mut tokens, config)?;
                 instructions.push(BasicAstInstruction::Assign(
                     variable_name.clone(),
                     expression,
@@ -285,11 +285,11 @@ pub fn build_ast(
                 if let Some((BasicToken::NewLine, _)) = tokens.get(0) {
                     instructions.push(BasicAstInstruction::Print(expressions));
                 } else {
-                    expressions.push((parse_expression(&mut tokens)?, false));
+                    expressions.push((parse_expression(&mut tokens, config)?, false));
 
                     while let Some((BasicToken::Comma, _)) = tokens.get(0) {
                         tokens.take(1);
-                        expressions.push((parse_expression(&mut tokens)?, false));
+                        expressions.push((parse_expression(&mut tokens, config)?, false));
                     }
 
                     instructions.push(BasicAstInstruction::Print(expressions));
@@ -307,7 +307,7 @@ pub fn build_ast(
                 let mut arguments = Vec::new();
 
                 while tokens.get(0).map(|pair| &pair.0) != Some(&BasicToken::CloseParen) {
-                    arguments.push(parse_expression(&mut tokens)?);
+                    arguments.push(parse_expression(&mut tokens, config)?);
 
                     match tokens.get(0) {
                         Some((BasicToken::Comma, _)) => {
@@ -440,6 +440,7 @@ impl_op_basic_ast_expression!(std::ops::Div, div, Operator::Div);
 
 pub(crate) fn parse_expression(
     tokens: &mut Cursor<'_, (BasicToken, Position)>,
+    config: &Config,
 ) -> Result<BasicAstExpression, ParseError> {
     /// Returns the first non-newline token in `tokens`
     fn peek(tokens: &[(BasicToken, Position)]) -> Option<&BasicToken> {
@@ -452,6 +453,7 @@ pub(crate) fn parse_expression(
     /// Parses a single expression item
     fn parse_expression_item(
         tokens: &mut Cursor<'_, (BasicToken, Position)>,
+        config: &Config,
     ) -> Result<BasicAstExpression, ParseError> {
         let position = tokens.get(0).map(|pair| pair.1).unwrap_or_default();
 
@@ -469,7 +471,7 @@ pub(crate) fn parse_expression(
                 let fn_name_lowercase = fn_name.to_ascii_lowercase();
                 let mut arguments = Vec::new();
                 while tokens.get(0).map(|pair| &pair.0) != Some(&BasicToken::CloseParen) {
-                    arguments.push(parse_expression(tokens)?);
+                    arguments.push(parse_expression(tokens, config)?);
 
                     match tokens.get(0).map(|pair| &pair.0) {
                         Some(BasicToken::Comma) => {
@@ -520,6 +522,14 @@ pub(crate) fn parse_expression(
                             Box::new(iter.next().unwrap()),
                         ))
                     }
+                } else if let Some(function_config) =
+                    config.builtin_functions.get(&fn_name_lowercase)
+                {
+                    function_config.validate_args(&arguments, span)?;
+                    Ok(BasicAstExpression::BuiltinFunction(
+                        fn_name_lowercase,
+                        arguments,
+                    ))
                 } else {
                     unimplemented!(
                         "User function calls are not yet supported! Function: {:?}",
@@ -537,7 +547,7 @@ pub(crate) fn parse_expression(
             }
             [BasicToken::OpenParen, ..] => {
                 tokens.take(1);
-                let res = parse_expression(tokens)?;
+                let res = parse_expression(tokens, config)?;
                 if let Some((BasicToken::CloseParen, _)) = tokens.take(1).get(0) {
                     Ok(res)
                 } else {
@@ -563,6 +573,7 @@ pub(crate) fn parse_expression(
         tokens: &mut Cursor<'_, (BasicToken, Position)>,
         lhs: BasicAstExpression,
         min_precedence: u8,
+        config: &Config,
     ) -> Result<BasicAstExpression, ParseError> {
         let mut ast = lhs;
         while let Some(&BasicToken::Operator(operator)) = peek(tokens) {
@@ -570,10 +581,10 @@ pub(crate) fn parse_expression(
                 break;
             }
             tokens.take(1);
-            let mut rhs = parse_expression_item(tokens)?;
+            let mut rhs = parse_expression_item(tokens, config)?;
             while let Some(&BasicToken::Operator(sub_operator)) = peek(tokens) {
                 if sub_operator.precedence() > operator.precedence() {
-                    rhs = parse_expression_main(tokens, rhs, operator.precedence() + 1)?;
+                    rhs = parse_expression_main(tokens, rhs, operator.precedence() + 1, config)?;
                 } else {
                     break;
                 }
@@ -586,8 +597,8 @@ pub(crate) fn parse_expression(
     }
 
     // Remove starting newlines
-    let lhs = parse_expression_item(tokens)?;
-    let res = parse_expression_main(tokens, lhs, 0)?;
+    let lhs = parse_expression_item(tokens, config)?;
+    let res = parse_expression_main(tokens, lhs, 0, config)?;
 
     Ok(res)
 }

src/parse/test.rs

@@ -130,22 +130,28 @@ fn test_parse_for() {
 #[test]
 fn test_operator_precedence() {
     fn test_parse<const N: usize>(list: [BasicToken; N]) -> BasicAstExpression {
-        parse_expression(&mut Cursor::from(
-            &list
-                .into_iter()
-                .map(|token| (token, Position::default()))
-                .collect::<Vec<_>>()[..],
-        ))
+        parse_expression(
+            &mut Cursor::from(
+                &list
+                    .into_iter()
+                    .map(|token| (token, Position::default()))
+                    .collect::<Vec<_>>()[..],
+            ),
+            &Config::default(),
+        )
         .unwrap()
     }
 
     fn test_err<const N: usize>(list: [BasicToken; N]) -> ParseErrorKind {
-        parse_expression(&mut Cursor::from(
-            &list
-                .into_iter()
-                .map(|token| (token, Position::default()))
-                .collect::<Vec<_>>()[..],
-        ))
+        parse_expression(
+            &mut Cursor::from(
+                &list
+                    .into_iter()
+                    .map(|token| (token, Position::default()))
+                    .collect::<Vec<_>>()[..],
+            ),
+            &Config::default(),
+        )
         .err()
         .unwrap()
         .kind

src/repr/basic.rs

@@ -12,6 +12,7 @@ pub enum BasicAstExpression {
         Box<BasicAstExpression>,
     ),
     Unary(UnaryOperator, Box<BasicAstExpression>),
+    BuiltinFunction(String, Vec<BasicAstExpression>),
 }
 
 #[derive(Clone, Debug, PartialEq)]

src/repr/operator.rs

@@ -24,6 +24,7 @@ pub enum Operator {
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 pub enum BasicOperator {
     Operator(Operator),
+    // TODO: rename to something else and have it not use variables as rhs keys
     Sensor,
 }
 

src/translate/mod.rs

@@ -40,6 +40,7 @@ fn translate_expression(
     expression: &BasicAstExpression,
     namer: &mut Namer,
     target_name: String,
+    config: &Config,
 ) -> MindustryProgram {
     match expression {
         BasicAstExpression::Integer(int) => {
@@ -62,8 +63,8 @@ fn translate_expression(
             let left_name = namer.temporary();
             let right_name = namer.temporary();
 
-            let mut res = translate_expression(left.as_ref(), namer, left_name.clone());
-            let mut right = translate_expression(right.as_ref(), namer, right_name.clone());
+            let mut res = translate_expression(left.as_ref(), namer, left_name.clone(), config);
+            let mut right = translate_expression(right.as_ref(), namer, right_name.clone(), config);
 
             res.append(&mut right);
 
@@ -89,7 +90,7 @@ fn translate_expression(
         }
         BasicAstExpression::Unary(op, value) => {
             let tmp_name = namer.temporary();
-            let mut res = translate_expression(value.as_ref(), namer, tmp_name.clone());
+            let mut res = translate_expression(value.as_ref(), namer, tmp_name.clone(), config);
 
             res.push(MindustryOperation::UnaryOperator(
                 target_name.clone(),
@@ -97,6 +98,30 @@ fn translate_expression(
                 Operand::Variable(tmp_name),
             ));
 
+            res
+        }
+        BasicAstExpression::BuiltinFunction(name, arguments) => {
+            let names = (0..arguments.len())
+                .map(|_| namer.temporary())
+                .collect::<Vec<_>>();
+
+            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
         }
     }
@@ -124,10 +149,16 @@ macro_rules! translate_operand {
     (
         $operand:expr,
         $res:expr,
-        $namer:expr $(,)?
+        $namer:expr,
+        $config:expr $(,)?
     ) => {{
         let name = $namer.temporary();
-        $res.append(&mut translate_expression($operand, $namer, name.clone()));
+        $res.append(&mut translate_expression(
+            $operand,
+            $namer,
+            name.clone(),
+            $config,
+        ));
         Operand::Variable(name)
     }};
 }
@@ -192,7 +223,8 @@ pub fn translate_ast(
                 res.push(MindustryOperation::End);
             }
             Instr::Assign(name, expression) => {
-                let mut instructions = translate_expression(expression, namer, name.clone());
+                let mut instructions =
+                    translate_expression(expression, namer, name.clone(), config);
                 res.append(&mut instructions);
             }
             Instr::IfThenElse(condition, true_branch, false_branch) => {
@@ -201,6 +233,7 @@ pub fn translate_ast(
                     condition,
                     namer,
                     condition_name.clone(),
+                    config,
                 ));
 
                 if !false_branch.instructions.is_empty() {
@@ -248,9 +281,24 @@ pub fn translate_ast(
                 let step_name = namer.temporary();
 
                 // Initialization: evaluate `start`, `end` and `step`
-                res.append(&mut translate_expression(start, namer, variable.clone()));
-                res.append(&mut translate_expression(end, namer, end_name.clone()));
-                res.append(&mut translate_expression(step, namer, step_name.clone()));
+                res.append(&mut translate_expression(
+                    start,
+                    namer,
+                    variable.clone(),
+                    config,
+                ));
+                res.append(&mut translate_expression(
+                    end,
+                    namer,
+                    end_name.clone(),
+                    config,
+                ));
+                res.append(&mut translate_expression(
+                    step,
+                    namer,
+                    step_name.clone(),
+                    config,
+                ));
 
                 // Condition
                 res.push(MindustryOperation::JumpLabel(start_label.clone()));
@@ -285,6 +333,7 @@ pub fn translate_ast(
                     condition,
                     namer,
                     condition_name.clone(),
+                    config,
                 ));
                 res.push(MindustryOperation::JumpIf(
                     end_label.clone(),
@@ -315,6 +364,7 @@ pub fn translate_ast(
                     condition,
                     namer,
                     condition_name.clone(),
+                    config,
                 ));
                 res.push(MindustryOperation::JumpIf(
                     start_label,
@@ -330,6 +380,7 @@ pub fn translate_ast(
                         &expression.0,
                         namer,
                         tmp_name.clone(),
+                        config,
                     ));
                     res.push(MindustryOperation::Print(Operand::Variable(tmp_name)));
                 }
@@ -341,8 +392,8 @@ pub fn translate_ast(
                 [BasicAstExpression::Variable(out_name), cell, index],
                 MindustryOperation::Read {
                     out_name: out_name.clone(),
-                    cell: translate_operand!(cell, res, namer),
-                    index: translate_operand!(index, res, namer),
+                    cell: translate_operand!(cell, res, namer, config),
+                    index: translate_operand!(index, res, namer, config),
                 }
             ),
             Instr::CallBuiltin(name, arguments) if name == "write" => translate_call!(
@@ -351,9 +402,9 @@ pub fn translate_ast(
                 res,
                 [value, cell, index],
                 MindustryOperation::Write {
-                    value: translate_operand!(value, res, namer),
-                    cell: translate_operand!(cell, res, namer),
-                    index: translate_operand!(index, res, namer),
+                    value: translate_operand!(value, res, namer, config),
+                    cell: translate_operand!(cell, res, namer, config),
+                    index: translate_operand!(index, res, namer, config),
                 }
             ),
             Instr::CallBuiltin(name, arguments) if name == "print_flush" => translate_call!(
@@ -361,7 +412,7 @@ pub fn translate_ast(
                 arguments,
                 res,
                 [cell],
-                MindustryOperation::PrintFlush(translate_operand!(cell, res, namer))
+                MindustryOperation::PrintFlush(translate_operand!(cell, res, namer, config))
             ),
             Instr::CallBuiltin(name, arguments) if name == "print_flush_global" => {
                 let BasicAstExpression::Variable(buffer) = &arguments[0] else {
@@ -371,12 +422,18 @@ pub fn translate_ast(
                 let instruction = MindustryOperation::WorldPrintFlush(match buffer.as_str() {
                     "mission" => WorldPrintFlush::Mission,
                     "notify" => WorldPrintFlush::Notify,
-                    "announce" => {
-                        WorldPrintFlush::Announce(translate_operand!(&arguments[1], res, namer))
-                    }
-                    "toast" => {
-                        WorldPrintFlush::Toast(translate_operand!(&arguments[1], res, namer))
-                    }
+                    "announce" => WorldPrintFlush::Announce(translate_operand!(
+                        &arguments[1],
+                        res,
+                        namer,
+                        config
+                    )),
+                    "toast" => WorldPrintFlush::Toast(translate_operand!(
+                        &arguments[1],
+                        res,
+                        namer,
+                        config
+                    )),
                     _ => unreachable!("print_flush_global constructed with invalid arguments"),
                 });
 
@@ -392,7 +449,7 @@ pub fn translate_ast(
                     arguments
                         .iter()
                         .skip(1)
-                        .map(|arg| translate_operand!(arg, res, namer))
+                        .map(|arg| translate_operand!(arg, res, namer, config))
                         .collect()
                 )
             ),
@@ -413,6 +470,7 @@ pub fn translate_ast(
                         argument,
                         namer,
                         argument_names[i].clone(),
+                        config,
                     ));
                 }
 

tests/examples/sensor.0.mlog

@@ -0,0 +1,3 @@
+set main__tmp_0 @unit
+set main__tmp_1 @health
+sensor health main__tmp_0 main__tmp_1

tests/examples/sensor.1.mlog

@@ -0,0 +1,2 @@
+set main__tmp_0 @unit
+sensor health main__tmp_0 @health