refactor: adjusted some macros in kclvm-parser and add some methods in kclvm-ast

kclvm/ast/src/ast.rs

@@ -42,7 +42,7 @@ use kclvm_span::Loc;
 use rustc_span::Pos;
 
 use super::token;
-
+use crate::node_ref;
 /// Node is the file, line and column number information
 /// that all AST nodes need to contain.
 #[derive(Serialize, Deserialize, Debug, Clone)]
@@ -225,6 +225,19 @@ pub struct Module {
     pub comments: Vec<NodeRef<Comment>>,
 }
 
+impl Module {
+    /// Get all ast.schema_stmts from ast.module and return it in a Vec.
+    pub fn filter_schema_stmt_from_module(&self) -> Vec<NodeRef<SchemaStmt>> {
+        let mut stmts = Vec::new();
+        for stmt in &self.body {
+            if let Stmt::Schema(schema_stmt) = &stmt.node {
+                stmts.push(node_ref!(schema_stmt.clone()));
+            }
+        }
+        return stmts;
+    }
+}
+
 /*
  * Stmt
  */
@@ -998,6 +1011,19 @@ pub struct StringLit {
     pub value: String,
 }
 
+/// Generate ast.StringLit from String
+impl TryFrom<String> for StringLit {
+    type Error = &'static str;
+
+    fn try_from(value: String) -> Result<Self, Self::Error> {
+        Ok(Self {
+            value: value.clone(),
+            raw_value: format!("{:?}", value),
+            is_long_string: false,
+        })
+    }
+}
+
 /// NameConstant, e.g.
 /// ```kcl
 /// True
@@ -1024,6 +1050,18 @@ impl NameConstant {
     }
 }
 
+/// Generate ast.NameConstant from Bool
+impl TryFrom<bool> for NameConstant {
+    type Error = &'static str;
+
+    fn try_from(value: bool) -> Result<Self, Self::Error> {
+        match value {
+            true => Ok(NameConstant::True),
+            false => Ok(NameConstant::False),
+        }
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug, Clone)]
 pub struct NameConstantLit {
     pub value: NameConstant,

kclvm/ast/src/lib.rs

@@ -1,4 +1,5 @@
 // Copyright 2021 The KCL Authors. All rights reserved.
+use crate::ast::*;
 
 pub mod ast;
 pub mod token;
@@ -9,3 +10,51 @@ pub mod walker;
 mod tests;
 
 pub const MAIN_PKG: &str = "__main__";
+
+#[macro_export]
+macro_rules! node_ref {
+    ($node: expr) => {
+        NodeRef::new(Node::dummy_node($node))
+    };
+    ($node: expr, $pos:expr) => {
+        NodeRef::new(Node::node_with_pos($node, $pos))
+    };
+}
+
+#[macro_export]
+macro_rules! expr_as {
+    ($expr: expr, $expr_enum: path) => {
+        if let $expr_enum(x) = ($expr.node as Expr) {
+            Some(x)
+        } else {
+            None
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! stmt_as {
+    ($stmt: expr, $stmt_enum: path) => {
+        if let $stmt_enum(x) = ($stmt.node as Stmt) {
+            Some(x)
+        } else {
+            None
+        }
+    };
+}
+
+/// Construct an AssignStmt node with assign_value as value
+pub fn build_assign_node(attr_name: &str, assign_value: NodeRef<Expr>) -> NodeRef<Stmt> {
+    let iden = node_ref!(Identifier {
+        names: vec![attr_name.to_string()],
+        pkgpath: String::new(),
+        ctx: ExprContext::Store
+    });
+
+    node_ref!(Stmt::Assign(AssignStmt {
+        value: assign_value,
+        targets: vec![iden],
+        type_annotation: None,
+        ty: None
+    }))
+}

kclvm/ast/src/tests.rs

@@ -1,5 +1,7 @@
-use crate::ast;
 use crate::walker::MutSelfMutWalker;
+use crate::AugOp::Mod;
+use crate::{ast, build_assign_node, node_ref, Identifier, Module, Node, NodeRef, SchemaStmt};
+use std::process::id;
 
 fn get_dummy_assign_ast() -> ast::Node<ast::AssignStmt> {
     let filename = "main.k";
@@ -140,3 +142,128 @@ fn test_mut_walker() {
     VarMutSelfMutWalker {}.walk_assign_stmt(&mut assign_stmt.node);
     assert_eq!(assign_stmt.node.targets[0].node.names[0], "x")
 }
+
+#[test]
+fn test_try_from_for_stringlit() {
+    let str_lit = ast::StringLit::try_from("test_str".to_string()).unwrap();
+    let json_str = serde_json::to_string(&str_lit).unwrap();
+
+    let str_expected =
+        r#"{"is_long_string":false,"raw_value":"\"test_str\"","value":"test_str"}"#
+            .to_string();
+    assert_eq!(str_expected, json_str);
+}
+
+#[test]
+fn test_try_from_for_nameconstant() {
+    let name_cons = ast::NameConstant::try_from(true).unwrap();
+    let json_str = serde_json::to_string(&name_cons).unwrap();
+    assert_eq!("\"True\"", json_str);
+
+    let name_cons = ast::NameConstant::try_from(false).unwrap();
+    let json_str = serde_json::to_string(&name_cons).unwrap();
+    assert_eq!("\"False\"", json_str);
+}
+
+#[test]
+fn test_filter_schema_with_no_schema() {
+    let ast_mod = Module {
+        filename: "".to_string(),
+        pkg: "".to_string(),
+        doc: "".to_string(),
+        name: "".to_string(),
+        body: vec![],
+        comments: vec![],
+    };
+    let schema_stmts = ast_mod.filter_schema_stmt_from_module();
+    assert_eq!(schema_stmts.len(), 0);
+}
+
+#[test]
+fn test_filter_schema_with_one_schema() {
+    let mut ast_mod = Module {
+        filename: "".to_string(),
+        pkg: "".to_string(),
+        doc: "".to_string(),
+        name: "".to_string(),
+        body: vec![],
+        comments: vec![],
+    };
+    let mut gen_schema_stmts = gen_schema_stmt(1);
+    ast_mod.body.append(&mut gen_schema_stmts);
+    let schema_stmts = ast_mod.filter_schema_stmt_from_module();
+    assert_eq!(schema_stmts.len(), 1);
+    assert_eq!(schema_stmts[0].node.name.node, "schema_stmt_0".to_string());
+}
+
+#[test]
+fn test_filter_schema_with_mult_schema() {
+    let mut ast_mod = Module {
+        filename: "".to_string(),
+        pkg: "".to_string(),
+        doc: "".to_string(),
+        name: "".to_string(),
+        body: vec![],
+        comments: vec![],
+    };
+    let mut gen_schema_stmts = gen_schema_stmt(10);
+    ast_mod.body.append(&mut gen_schema_stmts);
+    let schema_stmts = ast_mod.filter_schema_stmt_from_module();
+    assert_eq!(schema_stmts.len(), 10);
+    for i in 0..10 {
+        assert_eq!(
+            schema_stmts[i].node.name.node,
+            "schema_stmt_".to_string() + &i.to_string()
+        )
+    }
+}
+
+#[test]
+fn test_build_assign_stmt() {
+    let test_expr = node_ref!(ast::Expr::Identifier(Identifier {
+        names: vec!["name1".to_string(), "name2".to_string()],
+        pkgpath: "test".to_string(),
+        ctx: ast::ExprContext::Load
+    }));
+    let assgin_stmt = build_assign_node("test_attr_name", test_expr);
+
+    if let ast::Stmt::Assign(ref assign) = assgin_stmt.node {
+        if let ast::Expr::Identifier(ref iden) = &assign.value.node {
+            assert_eq!(iden.names.len(), 2);
+            assert_eq!(iden.names[0], "name1".to_string());
+            assert_eq!(iden.names[1], "name2".to_string());
+            assert_eq!(iden.pkgpath, "test".to_string());
+            match iden.ctx {
+                ast::ExprContext::Load => {}
+                _ => {
+                    assert!(false);
+                }
+            }
+        } else {
+            assert!(false);
+        }
+    } else {
+        assert!(false);
+    }
+}
+
+fn gen_schema_stmt(count: i32) -> Vec<NodeRef<ast::Stmt>> {
+    let mut schema_stmts = Vec::new();
+    for c in 0..count {
+        schema_stmts.push(node_ref!(ast::Stmt::Schema(SchemaStmt {
+            doc: "".to_string(),
+            name: node_ref!("schema_stmt_".to_string() + &c.to_string()),
+            parent_name: None,
+            for_host_name: None,
+            is_mixin: false,
+            is_protocol: false,
+            args: None,
+            mixins: vec![],
+            body: vec![],
+            decorators: vec![],
+            checks: vec![],
+            index_signature: None
+        })))
+    }
+    schema_stmts
+}

kclvm/parser/src/parser/expr.rs

@@ -8,8 +8,8 @@ use super::int::bytes_to_int;
 use super::Parser;
 
 use either::{self, Either};
+use kclvm_ast::node_ref;
 
-use crate::node_ref;
 use crate::parser::precedence::Precedence;
 use kclvm_ast::ast::*;
 use kclvm_ast::token;

kclvm/parser/src/parser/mod.rs

@@ -30,38 +30,6 @@ use kclvm_ast::token::{CommentKind, Token, TokenKind};
 use kclvm_ast::token_stream::{Cursor, TokenStream};
 use kclvm_span::symbol::Symbol;
 
-#[macro_export]
-macro_rules! node_ref {
-    ($node: expr) => {
-        NodeRef::new(Node::dummy_node($node))
-    };
-    ($node: expr, $pos:expr) => {
-        NodeRef::new(Node::node_with_pos($node, $pos))
-    };
-}
-
-#[macro_export]
-macro_rules! expr_as {
-    ($expr: expr, $expr_enum: path) => {
-        if let $expr_enum(x) = ($expr.node as Expr) {
-            Some(x)
-        } else {
-            None
-        }
-    };
-}
-
-#[macro_export]
-macro_rules! stmt_as {
-    ($stmt: expr, $stmt_enum: path) => {
-        if let $stmt_enum(x) = ($stmt.node as Stmt) {
-            Some(x)
-        } else {
-            None
-        }
-    };
-}
-
 pub struct Parser<'a> {
     /// The current token.
     pub token: Token,

kclvm/parser/src/parser/stmt.rs

@@ -3,15 +3,12 @@
 
 use core::panic;
 
-use kclvm_ast::ast::*;
+use kclvm_ast::{ast::*, node_ref, expr_as};
 use kclvm_ast::token::{DelimToken, LitKind, Token, TokenKind};
 use kclvm_span::symbol::kw;
 
 use super::Parser;
 
-use crate::expr_as;
-use crate::node_ref;
-
 /// Parser implementation of statements, which consists of expressions and tokens.
 /// Parser uses `parse_exprlist` and `parse_expr` in [`kclvm_parser::parser::expr`]
 /// to get a expression node, and then concretize it into the specified expression node,

kclvm/parser/src/parser/ty.rs

@@ -2,9 +2,7 @@
 
 use super::Parser;
 
-use crate::expr_as;
-
-use kclvm_ast::ast;
+use kclvm_ast::{ast, expr_as};
 use kclvm_ast::ast::{Expr, Node, NodeRef, Type};
 use kclvm_ast::token;
 use kclvm_ast::token::{BinOpToken, DelimToken, TokenKind};