From e569ce26b18e121aecee813d4448e07c8b93c834 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=9F=B3=E5=8D=9A=E6=96=87?= <sbw@sbw.so>
Date: Sat, 13 Jul 2024 14:07:30 +0800
Subject: [PATCH] feat: WIP: impl if statement

---
 lib/src/backend/lua/mod.rs                  | 38 +++++++++++++++++----
 lib/src/backend/lua/test.rs                 | 16 ++++-----
 viewer/test_projects/example1/test_proj.xml | 11 +++---
 3 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/lib/src/backend/lua/mod.rs b/lib/src/backend/lua/mod.rs
index 4c1147a..5f21187 100644
--- a/lib/src/backend/lua/mod.rs
+++ b/lib/src/backend/lua/mod.rs
@@ -11,8 +11,8 @@ use register::*;
 mod utils;
 use utils::*;
 
-mod vm;
 mod label;
+mod vm;
 use label::InstLabel;
 
 #[cfg(test)]
@@ -25,9 +25,9 @@ use crate::prelude::*;
 use indexmap::{IndexMap, IndexSet};
 use log::*;
 use smallvec::{smallvec, SmallVec};
+use std::collections::HashMap;
 use std::mem;
 use std::rc::Rc;
-use std::collections::HashMap;
 
 type ConstIdx = u8;
 
@@ -85,6 +85,7 @@ struct LuaBackendStates {
     scope: Option<Scope>,
     error: bool,
     access_mode: LuaAccessMode,
+    expr_exit_label: Option<InstLabel>,
 }
 
 impl LuaBackendStates {
@@ -108,6 +109,7 @@ impl Default for LuaBackendStates {
             error: false,
             access_mode: LuaAccessMode::None,
             const_idx: None,
+            expr_exit_label: None,
         }
     }
 }
@@ -159,6 +161,12 @@ impl LuaBackend {
         self.push_code(LuaByteCode::LoadK(r, k))
     }
 
+    #[inline]
+    fn code_jmp(&mut self, offset: i32) -> usize {
+        self.push_code(LuaByteCode::Jmp(offset));
+        self.byte_codes.len()
+    }
+
     #[inline]
     fn push_code(&mut self, code: LuaByteCode) {
         trace!("Code-Lua: {:?}", code);
@@ -241,6 +249,16 @@ impl LuaBackend {
         self.states.push(attr);
     }
 
+    fn push_exit_label(&mut self, exit_label: InstLabel) {
+        let attr = LuaBackendStates {
+            scope: self.top_attribute().scope.clone(),
+            expr_exit_label: Some(exit_label),
+            ..Default::default()
+        };
+
+        self.states.push(attr);
+    }
+
     #[inline]
     fn pop_attribute(&mut self) -> LuaBackendStates {
         self.states.pop().unwrap()
@@ -369,7 +387,10 @@ impl CodeGenBackend for LuaBackend {
     fn insert_label<S: AsRef<str>>(&mut self, label: S) {
         let label_name: StString = label.as_ref().into();
 
-        let label = self.labels.entry(label_name.clone()).or_insert(InstLabel::new(label_name));
+        let label = self
+            .labels
+            .entry(label_name.clone())
+            .or_insert(InstLabel::new(label_name));
 
         // ensure label is not inserted
         debug_assert!(label.inst_index.is_some());
@@ -513,10 +534,9 @@ impl AstVisitorMut for LuaBackend {
     fn visit_if_statement_mut(&mut self, ifst: &mut IfStatement) {
         trace!("LuaGen: if statement: {}", ifst.condition());
 
-        let cond_true = self.create_label("if-true");
-        let cond_false = self.create_label("if-false");
+        let if_exit_label = self.create_label("if-exit");
 
-        self.push_default_attribute();
+        self.push_exit_label(if_exit_label);
         self.visit_expression_mut(ifst.condition_mut());
         let attr = self.pop_attribute();
 
@@ -546,6 +566,7 @@ impl AstVisitorMut for LuaBackend {
                     Some(r) => *r,
                     _ => self.reg_mgr.alloc_hard(),
                 };
+                let cond = self.top_attribute().expr_exit_label.clone();
 
                 self.push_access_attribute(LuaAccessMode::LoadNewRegister);
                 self.visit_expression_mut(&mut operands[0]);
@@ -564,6 +585,11 @@ impl AstVisitorMut for LuaBackend {
                     _ => unreachable!(),
                 }
 
+                // conditional jump
+                if let Some(lbl) = cond {
+                    let fixup = self.code_jmp(0);
+                }
+
                 if let RK::R(r0) = rk0 {
                     self.reg_mgr.free(&r0);
                 }
diff --git a/lib/src/backend/lua/test.rs b/lib/src/backend/lua/test.rs
index e5360a8..a42955c 100644
--- a/lib/src/backend/lua/test.rs
+++ b/lib/src/backend/lua/test.rs
@@ -1,9 +1,9 @@
+use mlua::{ChunkMode, Lua};
 use std::io::Write;
 use std::process::Command;
-use mlua::{ChunkMode, Lua};
 
-use crate::{parser::*, prelude::*,};
 use crate::backend::{CodeGenBackend, CodeGenDriver, LuaBackend};
+use crate::{parser::*, prelude::*};
 
 fn generate_module<S1: AsRef<str>, S2: AsRef<str>>(decl: S1, body: S2, writer: &mut dyn Write) {
     let mgr = UnitsManager::new();
@@ -119,7 +119,7 @@ fn test_add() {
 
     let lua = Lua::new();
     assert!(lua.load(buf).set_mode(ChunkMode::Binary).exec().is_ok());
-    
+
     let r = lua.globals().get::<_, i32>("a");
     assert_eq!(r.unwrap(), 3);
 }
@@ -133,9 +133,9 @@ fn test_if_statement() {
     let mut buf = vec![];
     generate_module(decl, body, &mut buf);
 
-    // let lua = Lua::new();
-    // assert!(lua.load(buf).set_mode(ChunkMode::Binary).exec().is_ok());
-    // 
-    // let r = lua.globals().get::<_, i32>("a");
-    // assert_eq!(r.unwrap(), 0);
+    let lua = Lua::new();
+    assert!(lua.load(buf).set_mode(ChunkMode::Binary).exec().is_ok());
+
+    let r = lua.globals().get::<_, i32>("a");
+    assert_eq!(r.unwrap(), 0);
 }
diff --git a/viewer/test_projects/example1/test_proj.xml b/viewer/test_projects/example1/test_proj.xml
index 470d33f..3f975ed 100644
--- a/viewer/test_projects/example1/test_proj.xml
+++ b/viewer/test_projects/example1/test_proj.xml
@@ -18,11 +18,12 @@ END_VAR
 end_program
 ]]>
             </interface>
-            <body>
-                a := 1;
-                if A = 2 then
-                a := 0;
-                end_if
+            <body><![CDATA[
+a := 1;
+if A = 2 then
+    a := 0;
+end_if
+]]>
             </body>
         </pou>
     </pou-list>