1  /* This file is part of DarkFi (https://dark.fi)
  2   *
  3   * Copyright (C) 2020-2025 Dyne.org foundation
  4   *
  5   * This program is free software: you can redistribute it and/or modify
  6   * it under the terms of the GNU Affero General Public License as
  7   * published by the Free Software Foundation, either version 3 of the
  8   * License, or (at your option) any later version.
  9   *
 10   * This program is distributed in the hope that it will be useful,
 11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
 12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 13   * GNU Affero General Public License for more details.
 14   *
 15   * You should have received a copy of the GNU Affero General Public License
 16   * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 17   */
 18  
 19  use crate::{
 20      error::{Error, Result},
 21      //prop::{Property, PropertySubType, PropertyType, PropertySExprValue},
 22  };
 23  use std::collections::HashMap;
 24  
 25  use super::{Op, SExprCode};
 26  
 27  #[derive(Debug, Clone)]
 28  enum Token {
 29      LoadVar(String),
 30      Add,
 31      Sub,
 32      Mul,
 33      Div,
 34      LeftParen,
 35      RightParen,
 36      ConstFloat32(f32),
 37      NestedExpr(Box<Vec<Token>>),
 38      SubExpr(Box<Vec<Token>>),
 39      If,
 40      Else,
 41      LeftBrace,
 42      RightBrace,
 43      LessThan,
 44      IfElse((Box<Vec<Token>>, Box<Vec<Token>>, Box<Vec<Token>>)),
 45      LessThanCompare((Box<Vec<Token>>, Box<Vec<Token>>)),
 46      Equals,
 47      SetValue((String, Box<Vec<Token>>)),
 48  }
 49  
 50  impl Token {
 51      fn flatten(self) -> Vec<Self> {
 52          match self {
 53              Self::NestedExpr(tokens) => {
 54                  let tokens: Vec<_> = *tokens;
 55                  tokens.into_iter().map(|t| t.flatten()).flatten().collect()
 56              }
 57              _ => vec![self],
 58          }
 59      }
 60  }
 61  
 62  #[derive(Clone)]
 63  pub struct Compiler {
 64      table: HashMap<String, Token>,
 65  }
 66  
 67  impl Compiler {
 68      pub fn new() -> Self {
 69          Self { table: HashMap::new() }
 70      }
 71  
 72      pub fn add_const_f32<S: Into<String>>(&mut self, name: S, val: f32) {
 73          self.table.insert(name.into(), Token::ConstFloat32(val));
 74      }
 75      /*
 76      pub fn get_const_f32<S: AsRef<str>>(&self, name: S) -> Option<f32> {
 77          let name = name.as_ref();
 78          let val = self.table.get(name)?;
 79          match val {
 80              Token::ConstFloat32(v) => Some(v)
 81              _ => None
 82          }
 83      }
 84      */
 85  
 86      pub fn compile<S: AsRef<str>>(&self, prestmts: S) -> Result<SExprCode> {
 87          let prestmts = prestmts.as_ref();
 88          // Strip all comments
 89          let mut stmts = String::new();
 90          for line in prestmts.lines() {
 91              if let Some(chr) = line.trim_start().chars().next() {
 92                  if chr != '#' {
 93                      stmts.push_str(line);
 94                  }
 95              }
 96          }
 97  
 98          let mut code = vec![];
 99          for stmt in stmts.split(';') {
100              code.push(self.compile_line(stmt)?);
101          }
102          Ok(code)
103      }
104  
105      fn compile_line(&self, stmt: &str) -> Result<Op> {
106          let tokens = self.tokenize(&stmt);
107          //println!("{tokens:#?}");
108          let tokens = to_rpn(tokens)?;
109          //println!("{tokens:#?}");
110          Ok(convert(&mut tokens.into_iter())?)
111      }
112  
113      fn tokenize(&self, stmt: &str) -> Vec<Token> {
114          let mut tokens = Vec::new();
115          let mut current_token = String::new();
116          for chr in stmt.chars() {
117              match chr {
118                  ' ' | '\t' | '\n' => {
119                      self.clear_accum(&mut current_token, &mut tokens);
120                  }
121                  '+' => {
122                      self.clear_accum(&mut current_token, &mut tokens);
123                      tokens.push(Token::Add);
124                  }
125                  '-' => {
126                      self.clear_accum(&mut current_token, &mut tokens);
127                      tokens.push(Token::Sub);
128                  }
129                  '*' => {
130                      self.clear_accum(&mut current_token, &mut tokens);
131                      tokens.push(Token::Mul);
132                  }
133                  '/' => {
134                      self.clear_accum(&mut current_token, &mut tokens);
135                      tokens.push(Token::Div);
136                  }
137                  '(' => {
138                      self.clear_accum(&mut current_token, &mut tokens);
139                      tokens.push(Token::LeftParen);
140                  }
141                  ')' => {
142                      self.clear_accum(&mut current_token, &mut tokens);
143                      tokens.push(Token::RightParen);
144                  }
145                  '{' => {
146                      self.clear_accum(&mut current_token, &mut tokens);
147                      tokens.push(Token::LeftBrace);
148                  }
149                  '}' => {
150                      self.clear_accum(&mut current_token, &mut tokens);
151                      tokens.push(Token::RightBrace);
152                  }
153                  '<' => {
154                      self.clear_accum(&mut current_token, &mut tokens);
155                      tokens.push(Token::LessThan);
156                  }
157                  '=' => {
158                      self.clear_accum(&mut current_token, &mut tokens);
159                      tokens.push(Token::Equals);
160                  }
161                  _ => current_token.push(chr),
162              }
163          }
164          self.clear_accum(&mut current_token, &mut tokens);
165          tokens
166      }
167  
168      fn clear_accum(&self, current_token: &mut String, tokens: &mut Vec<Token>) {
169          let prev_token = std::mem::replace(current_token, String::new());
170          if prev_token.is_empty() {
171              return
172          }
173  
174          if let Some(token) = self.table.get(&prev_token) {
175              tokens.push(token.clone());
176              return
177          }
178  
179          match prev_token.as_str() {
180              "if" => {
181                  tokens.push(Token::If);
182                  return
183              }
184              "else" => {
185                  tokens.push(Token::Else);
186                  return
187              }
188              _ => {}
189          }
190  
191          // Number or var?
192          match prev_token.parse::<f32>() {
193              Ok(v) => tokens.push(Token::ConstFloat32(v)),
194              Err(_) => tokens.push(Token::LoadVar(prev_token)),
195          }
196      }
197  }
198  
199  /// Convert from infix to reverse polish notation
200  fn to_rpn(tokens: Vec<Token>) -> Result<Vec<Token>> {
201      //println!("to_rpn = {tokens:#?}");
202      let mut out;
203      let mut stack = Vec::new();
204  
205      // equals
206      let mut iter = tokens.into_iter();
207      let mut var = String::new();
208      let mut lhs = vec![];
209      let mut rhs = vec![];
210      let mut comparison = 0;
211      // 0: none
212      // 1: =
213      while let Some(token) = iter.next() {
214          match token {
215              Token::Equals => {
216                  if comparison != 0 {
217                      return Err(Error::UnexpectedToken)
218                  }
219                  comparison = 1;
220                  if lhs.len() != 1 {
221                      return Err(Error::UnexpectedToken)
222                  }
223                  let lhs = std::mem::take(&mut lhs);
224                  match lhs.into_iter().next().unwrap() {
225                      Token::LoadVar(v) => var = v,
226                      _ => return Err(Error::UnexpectedToken),
227                  }
228              }
229              token => {
230                  if comparison == 0 {
231                      lhs.push(token);
232                  } else {
233                      rhs.push(token);
234                  }
235              }
236          }
237      }
238      if comparison == 1 {
239          let stack = std::mem::take(&mut rhs);
240          let rpn = to_rpn(stack)?;
241          out = vec![Token::SetValue((var, Box::new(rpn)))];
242      } else {
243          assert!(rhs.is_empty());
244          out = lhs;
245      }
246  
247      // Parens
248      let tokens = std::mem::take(&mut out);
249      let mut paren = 0;
250      for token in tokens {
251          match token {
252              Token::LeftParen => {
253                  // Is this the first opening paren for this subexpr?
254                  if paren > 0 {
255                      stack.push(token);
256                  }
257                  paren += 1;
258                  continue
259              }
260              Token::RightParen => {
261                  paren -= 1;
262  
263                  // Whoops non-matching number of parens!
264                  if paren < 0 {
265                      return Err(Error::UnexpectedToken)
266                  }
267  
268                  // Did we finally reach the closing paren for this subexpr?
269                  if paren == 0 {
270                      let stack = std::mem::take(&mut stack);
271                      let rpn = to_rpn(stack)?;
272                      out.push(Token::NestedExpr(Box::new(rpn)));
273                  } else {
274                      stack.push(token);
275                  }
276  
277                  continue
278              }
279              _ => {}
280          }
281          if paren > 0 {
282              stack.push(token);
283          } else {
284              out.push(token);
285          }
286      }
287      out.append(&mut stack);
288  
289      // Braces
290      let tokens = std::mem::take(&mut out);
291      assert!(stack.is_empty());
292      let mut paren = 0;
293      for token in tokens {
294          match token {
295              Token::LeftBrace => {
296                  // Is this the first opening paren for this subexpr?
297                  if paren > 0 {
298                      stack.push(token);
299                  }
300                  paren += 1;
301                  continue
302              }
303              Token::RightBrace => {
304                  paren -= 1;
305  
306                  // Whoops non-matching number of parens!
307                  if paren < 0 {
308                      return Err(Error::UnexpectedToken)
309                  }
310  
311                  // Did we finally reach the closing paren for this subexpr?
312                  if paren == 0 {
313                      let stack = std::mem::take(&mut stack);
314                      let rpn = to_rpn(stack)?;
315                      out.push(Token::SubExpr(Box::new(rpn)));
316                  } else {
317                      stack.push(token);
318                  }
319  
320                  continue
321              }
322              _ => {}
323          }
324          if paren > 0 {
325              stack.push(token);
326          } else {
327              out.push(token);
328          }
329      }
330      out.append(&mut stack);
331  
332      let tokens = std::mem::take(&mut out);
333      assert!(stack.is_empty());
334      let mut iter = tokens.into_iter();
335      'mainloop: while let Some(token) = iter.next() {
336          match token {
337              Token::If => {
338                  let mut section = 0;
339                  let mut cond_expr = vec![];
340                  let mut if_expr = vec![];
341                  let mut else_expr = vec![];
342                  while let Some(token) = iter.next() {
343                      match token {
344                          Token::SubExpr(tokens) => {
345                              if section == 0 {
346                                  let cexpr = std::mem::take(&mut cond_expr);
347                                  cond_expr = to_rpn(cexpr)?;
348  
349                                  if_expr = *tokens;
350                              } else if section == 1 {
351                                  else_expr = *tokens;
352                              } else {
353                                  return Err(Error::UnexpectedToken)
354                              }
355  
356                              section += 1;
357                          }
358                          Token::Else => {
359                              if section != 1 {
360                                  return Err(Error::UnexpectedToken)
361                              }
362                          }
363                          token => {
364                              if section != 0 {
365                                  out.push(Token::IfElse((
366                                      Box::new(cond_expr),
367                                      Box::new(if_expr),
368                                      Box::new(else_expr),
369                                  )));
370                                  out.push(token);
371                                  continue 'mainloop;
372                              }
373                              cond_expr.push(token);
374                          }
375                      }
376                  }
377                  // We reached the end
378                  out.push(Token::IfElse((
379                      Box::new(cond_expr),
380                      Box::new(if_expr),
381                      Box::new(else_expr),
382                  )));
383              }
384              token => out.push(token),
385          }
386      }
387  
388      // comparisons <>=
389      let tokens = std::mem::take(&mut out);
390      let mut iter = tokens.into_iter();
391      let mut lhs = vec![];
392      let mut rhs = vec![];
393      let mut comparison = 0;
394      // 0: none
395      // 1: <
396      while let Some(token) = iter.next() {
397          match token {
398              Token::LessThan => {
399                  if comparison != 0 {
400                      return Err(Error::UnexpectedToken)
401                  }
402                  comparison = 1;
403              }
404              token => {
405                  if comparison == 0 {
406                      lhs.push(token);
407                  } else {
408                      rhs.push(token);
409                  }
410              }
411          }
412      }
413      if comparison == 1 {
414          out = vec![Token::LessThanCompare((Box::new(lhs), Box::new(rhs)))];
415      } else {
416          assert!(rhs.is_empty());
417          out = lhs;
418      }
419  
420      // */
421      let tokens = std::mem::take(&mut out);
422      assert!(stack.is_empty());
423      let mut is_op = false;
424      for token in tokens {
425          match token {
426              Token::Mul | Token::Div => {
427                  if is_op {
428                      return Err(Error::UnexpectedToken)
429                  }
430                  is_op = true;
431  
432                  let Some(prev_item) = out.pop() else { return Err(Error::UnexpectedToken) };
433                  stack.push(token);
434                  stack.push(prev_item);
435              }
436              _ => {
437                  if is_op {
438                      is_op = false;
439                      let mut expr = std::mem::take(&mut stack);
440                      expr.push(token);
441                      out.push(Token::NestedExpr(Box::new(expr)));
442  
443                      continue
444                  }
445  
446                  out.push(token)
447              }
448          }
449      }
450      //println!("out = {out:#?}");
451      //println!("stack = {stack:#?}");
452      //assert!(!is_op);
453      //assert!(stack.is_empty());
454      if is_op || !stack.is_empty() {
455          return Err(Error::UnexpectedToken)
456      }
457  
458      // +-
459      let tokens = std::mem::take(&mut out);
460      let mut is_op = false;
461      for token in tokens {
462          match token {
463              Token::Add | Token::Sub => {
464                  if is_op {
465                      return Err(Error::UnexpectedToken)
466                  }
467                  is_op = true;
468  
469                  let Some(prev_item) = out.pop() else { return Err(Error::UnexpectedToken) };
470                  stack.push(token);
471                  stack.push(prev_item);
472              }
473              _ => {
474                  if is_op {
475                      is_op = false;
476                      let mut expr = std::mem::take(&mut stack);
477                      expr.push(token);
478                      out.push(Token::NestedExpr(Box::new(expr)));
479  
480                      continue
481                  }
482  
483                  out.push(token)
484              }
485          }
486      }
487      //assert!(!is_op);
488      //assert!(stack.is_empty());
489      if is_op || !stack.is_empty() {
490          return Err(Error::UnexpectedToken)
491      }
492  
493      // Flatten everything
494      let out = out.into_iter().map(|t| t.flatten()).flatten().collect();
495      Ok(out)
496  }
497  
498  fn convert<I: Iterator<Item = Token>>(iter: &mut I) -> Result<Op> {
499      let Some(token) = iter.next() else { return Err(Error::UnexpectedToken) };
500  
501      let op = match token {
502          Token::ConstFloat32(v) => Op::ConstFloat32(v),
503          Token::LoadVar(v) => Op::LoadVar(v),
504          Token::Add => {
505              let lhs = convert(iter)?;
506              let rhs = convert(iter)?;
507              Op::Add((Box::new(lhs), Box::new(rhs)))
508          }
509          Token::Sub => {
510              let lhs = convert(iter)?;
511              let rhs = convert(iter)?;
512              Op::Sub((Box::new(lhs), Box::new(rhs)))
513          }
514          Token::Mul => {
515              let lhs = convert(iter)?;
516              let rhs = convert(iter)?;
517              Op::Mul((Box::new(lhs), Box::new(rhs)))
518          }
519          Token::Div => {
520              let lhs = convert(iter)?;
521              let rhs = convert(iter)?;
522              Op::Div((Box::new(lhs), Box::new(rhs)))
523          }
524          Token::IfElse((cond, if_val, else_val)) => {
525              let cond = convert(&mut cond.into_iter())?;
526              let if_val = convert(&mut if_val.into_iter())?;
527              let else_val = convert(&mut else_val.into_iter())?;
528              Op::IfElse((Box::new(cond), vec![if_val], vec![else_val]))
529          }
530          Token::LessThanCompare((lhs, rhs)) => {
531              let lhs = convert(&mut lhs.into_iter())?;
532              let rhs = convert(&mut rhs.into_iter())?;
533              Op::LessThan((Box::new(lhs), Box::new(rhs)))
534          }
535          Token::SetValue((var, expr)) => {
536              let expr = convert(&mut expr.into_iter())?;
537              Op::StoreVar((var, Box::new(expr)))
538          }
539          _ => return Err(Error::UnexpectedToken),
540      };
541      Ok(op)
542  }
543  
544  #[cfg(test)]
545  mod tests {
546      use super::*;
547  
548      #[test]
549      fn single_line() {
550          let compiler = Compiler::new();
551  
552          let code = compiler.compile("h/2 - 200").unwrap();
553          #[rustfmt::skip]
554          let code2 = vec![Op::Sub((
555              Box::new(Op::Div((
556                  Box::new(Op::LoadVar("h".to_string())),
557                  Box::new(Op::ConstFloat32(2.)),
558              ))),
559              Box::new(Op::ConstFloat32(200.)),
560          ))];
561          assert_eq!(code, code2);
562  
563          let code = compiler.compile("(x + h/2 + (y + 7)/5) - 200").unwrap();
564          #[rustfmt::skip]
565          let code2 = vec![Op::Sub((
566              Box::new(Op::Add((
567                  Box::new(Op::Add((
568                      Box::new(Op::LoadVar("x".to_string())),
569                      Box::new(Op::Div((
570                          Box::new(Op::LoadVar("h".to_string())),
571                          Box::new(Op::ConstFloat32(2.))
572                      ))),
573                  ))),
574  
575                  Box::new(Op::Div((
576                      Box::new(Op::Add((
577                          Box::new(Op::LoadVar("y".to_string())),
578                          Box::new(Op::ConstFloat32(7.))
579                      ))),
580                      Box::new(Op::ConstFloat32(5.))
581                  ))),
582              ))),
583  
584              Box::new(Op::ConstFloat32(200.))
585          ))];
586          assert_eq!(code, code2);
587      }
588  
589      #[test]
590      fn h_minus_1() {
591          let mut compiler = Compiler::new();
592          let code = compiler.compile("h - 1").unwrap();
593          #[rustfmt::skip]
594          let code2 = vec![Op::Sub((
595              Box::new(Op::LoadVar("h".to_string())),
596              Box::new(Op::ConstFloat32(1.))
597          ))];
598          assert_eq!(code, code2);
599      }
600  
601      #[test]
602      fn dosub() {
603          let mut compiler = Compiler::new();
604          compiler.add_const_f32("HELLO", 110.);
605  
606          let code = compiler.compile("HELLO").unwrap();
607          let code2 = vec![Op::ConstFloat32(110.)];
608          assert_eq!(code, code2);
609      }
610  
611      #[test]
612      fn if_else() {
613          let mut compiler = Compiler::new();
614          let code = compiler
615              .compile(
616                  "
617              if h < 4 {
618                  h - 1
619              } else {
620                  2 * h + 5
621              }
622          ",
623              )
624              .unwrap();
625          let code2 = vec![Op::IfElse((
626              Box::new(Op::LessThan((
627                  Box::new(Op::LoadVar("h".to_string())),
628                  Box::new(Op::ConstFloat32(4.)),
629              ))),
630              vec![Op::Sub((Box::new(Op::LoadVar("h".to_string())), Box::new(Op::ConstFloat32(1.))))],
631              vec![Op::Add((
632                  Box::new(Op::Mul((
633                      Box::new(Op::ConstFloat32(2.)),
634                      Box::new(Op::LoadVar("h".to_string())),
635                  ))),
636                  Box::new(Op::ConstFloat32(5.)),
637              ))],
638          ))];
639          assert_eq!(code, code2);
640      }
641  
642      #[test]
643      fn set_val() {
644          let mut compiler = Compiler::new();
645          let code = compiler
646              .compile(
647                  "
648              r = 10 / 4
649          ",
650              )
651              .unwrap();
652          let code2 = vec![Op::StoreVar((
653              "r".to_string(),
654              Box::new(Op::Div((Box::new(Op::ConstFloat32(10.)), Box::new(Op::ConstFloat32(4.))))),
655          ))];
656          assert_eq!(code, code2);
657      }
658  
659      #[test]
660      fn multiline() {
661          let mut compiler = Compiler::new();
662          let code = compiler
663              .compile(
664                  "
665              # This is a comment
666              r = 10 / 4;
667  
668              s = if h < 4 {
669                  h - 1
670              } else {
671                  2 * h + 5
672              };
673  
674              r + 1
675          ",
676              )
677              .unwrap();
678          let code2 = vec![
679              Op::StoreVar((
680                  "r".to_string(),
681                  Box::new(Op::Div((
682                      Box::new(Op::ConstFloat32(10.)),
683                      Box::new(Op::ConstFloat32(4.)),
684                  ))),
685              )),
686              Op::StoreVar((
687                  "s".to_string(),
688                  Box::new(Op::IfElse((
689                      Box::new(Op::LessThan((
690                          Box::new(Op::LoadVar("h".to_string())),
691                          Box::new(Op::ConstFloat32(4.)),
692                      ))),
693                      vec![Op::Sub((
694                          Box::new(Op::LoadVar("h".to_string())),
695                          Box::new(Op::ConstFloat32(1.)),
696                      ))],
697                      vec![Op::Add((
698                          Box::new(Op::Mul((
699                              Box::new(Op::ConstFloat32(2.)),
700                              Box::new(Op::LoadVar("h".to_string())),
701                          ))),
702                          Box::new(Op::ConstFloat32(5.)),
703                      ))],
704                  ))),
705              )),
706              Op::Add((Box::new(Op::LoadVar("r".to_string())), Box::new(Op::ConstFloat32(1.)))),
707          ];
708          assert_eq!(code, code2);
709      }
710  }