Merge remote-tracking branch 'origin/main' into paultag/refgraph

src/wasm-lib/Cargo.toml

@@ -9,6 +9,12 @@ rust-version = "1.73"
 [lib]
 crate-type = ["cdylib"]
 
+# DO NOT MERGE ME WITH THIS PLEASE
+#
+# FOR THE LOVE OF GOD DON'T HIT APPROVE PR WITHOUT ASKING ME TO REMOVE THIS
+[package.metadata.wasm-pack.profile.release]
+wasm-opt = false
+
 [dependencies]
 bson = { version = "2.13.0", features = ["uuid-1", "chrono"] }
 data-encoding = "2.6.0"

src/wasm-lib/kcl/src/execution/mod.rs

@@ -1,6 +1,6 @@
 //! The executor for the AST.
 
-use std::{path::PathBuf, sync::Arc};
+use std::{cmp::Ordering, path::PathBuf, sync::Arc};
 
 use annotations::AnnotationScope;
 use anyhow::Result;
@@ -43,6 +43,7 @@ use crate::{
         BodyItem, Expr, FunctionExpression, ImportSelector, ItemVisibility, Node, NodeRef, NonCodeValue,
         Program as AstProgram, TagDeclarator, TagNode,
     },
+    refgraph::extract_refgraph,
     settings::types::UnitLength,
     source_range::{ModuleId, SourceRange},
     std::{args::Arg, StdLib},

src/wasm-lib/kcl/src/lib.rs

@@ -66,6 +66,7 @@ pub mod lint;
 mod log;
 mod lsp;
 mod parsing;
+mod refgraph;
 mod settings;
 #[cfg(test)]
 mod simulation_tests;

src/wasm-lib/kcl/src/refgraph/mod.rs

@@ -0,0 +1,296 @@
+mod walk;
+
+use crate::{
+    parsing::ast::types,
+    walk::{Node, Visitable},
+    FormatOptions,
+};
+use std::sync::{Arc, Mutex};
+
+///
+type Declaration<'tree> = (Vec<Node<'tree>>, &'tree types::Identifier);
+
+///
+type Reference<'tree> = (Vec<Node<'tree>>, &'tree types::Identifier);
+
+///
+type RefEdge<'tree> = (Option<Declaration<'tree>>, Reference<'tree>);
+
+///
+#[derive(Clone, Debug)]
+pub struct Scope<'tree> {
+    ///
+    pub declarations: Vec<Declaration<'tree>>,
+
+    ///
+    pub references: Vec<Reference<'tree>>,
+
+    ///
+    pub children: Vec<Scope<'tree>>,
+}
+
+impl<'tree> Scope<'tree> {
+    pub fn new() -> Self {
+        Scope {
+            declarations: vec![],
+            references: vec![],
+            children: vec![],
+        }
+    }
+
+    pub fn edges(&self) -> Vec<RefEdge> {
+        let mut edges = Vec::<RefEdge>::new();
+        let mut unmatched_refs = self.references.clone();
+
+        for child in &self.children {
+            for (declaration, (ref_node, ref_id)) in child.edges() {
+                match declaration {
+                    Some((decl_node, decl_id)) => {
+                        edges.push((Some((decl_node.clone(), decl_id)), (ref_node.clone(), ref_id)));
+                    }
+                    None => {
+                        unmatched_refs.push((ref_node.clone(), ref_id));
+                    }
+                }
+            }
+        }
+
+        for (ref_node, ref_id) in unmatched_refs.into_iter() {
+            edges.push((
+                self.declarations
+                    .iter()
+                    .filter(|(_, decl_id)| decl_id.name == ref_id.name)
+                    .cloned()
+                    .next(),
+                (ref_node.clone(), ref_id),
+            ));
+        }
+
+        edges
+    }
+}
+
+///
+#[derive(Clone, Debug)]
+pub struct ScopeVisitor<'tree> {
+    scope: Arc<Mutex<Scope<'tree>>>,
+}
+
+impl<'tree> ScopeVisitor<'tree> {
+    pub fn new() -> Self {
+        Self {
+            scope: Arc::new(Mutex::new(Scope::new())),
+        }
+    }
+}
+
+impl<'tree> ScopeVisitor<'tree> {
+    fn visit(&self, path: Vec<Node<'tree>>, node: Node<'tree>) -> Result<bool, std::convert::Infallible> {
+        if let Node::Program(program) = node {
+            let csv = ScopeVisitor::new();
+            for child in node.children() {
+                let mut path = path.clone();
+                path.push(node.clone());
+                csv.clone().visit(path, child.clone())?;
+            }
+
+            self.scope
+                .lock()
+                .unwrap()
+                .children
+                .push(csv.scope.lock().unwrap().clone());
+            return Ok(true);
+        }
+
+        let mut path = path.clone();
+        path.push(node.clone());
+
+        match node {
+            Node::VariableDeclaration(vd) => {
+                self.scope
+                    .lock()
+                    .unwrap()
+                    .declarations
+                    .push((vec![node.clone()], &vd.declaration.id));
+
+                let node: Node = (&vd.declaration.init).into();
+                self.clone().visit(path, node)?;
+            }
+            Node::Identifier(id) => {
+                self.scope.lock().unwrap().references.push((vec![node.clone()], &id));
+            }
+            _ => {
+                for child in node.children() {
+                    self.visit(path.clone(), child.clone())?;
+                }
+            }
+        }
+
+        Ok(true)
+    }
+}
+
+pub fn extract_refgraph<'a, 'tree>(
+    program: types::NodeRef<'tree, types::Program>,
+) -> Result<Scope<'tree>, std::convert::Infallible> {
+    let sv = ScopeVisitor::new();
+    let node: Node = (program).into();
+
+    for child in node.children() {
+        sv.clone().visit(vec![node.clone()], child.clone())?;
+    }
+
+    let x = sv.scope.lock().unwrap().clone();
+    Ok(x)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    macro_rules! kcl {
+        ( $kcl:expr ) => {{
+            $crate::parsing::top_level_parse($kcl).unwrap()
+        }};
+    }
+
+    #[test]
+    fn extract_edges() {
+        let program = kcl!(
+            "
+const foo = 1
+const bar = foo + 1
+
+fn myfn = () => {
+    const foo = 2
+    sin(foo)
+}
+"
+        );
+
+        let refgraph = extract_refgraph(&program).unwrap();
+
+        let edges = refgraph.edges().into_iter().collect::<Vec<_>>();
+        assert_eq!(edges.len(), 3);
+
+        // sin is a global so we are chilin
+        let (decl, (_ref_node, ref_id)) = edges.get(2).unwrap();
+        assert_eq!("sin", ref_id.name);
+        assert!(decl.is_none());
+
+        // myfn's foo
+        let (decl, (_ref_node, ref_id)) = edges.get(1).unwrap();
+        assert_eq!("foo", ref_id.name);
+        assert!(decl.is_some());
+        // todo: check this is actually refering to the parent scope foo
+
+        let (Some((decl_nodes, _)), _) = edges.get(1).unwrap() else {
+            panic!("not some");
+        };
+
+        assert!(decl_nodes[0].ptr_eq((&program.body[0]).into()));
+    }
+
+    #[test]
+    fn extract_simple() {
+        let program = kcl!(
+            "
+const foo = 1
+const bar = foo + 1
+
+fn myfn = () => {
+    const foo = 2
+    sin(foo)
+}
+"
+        );
+
+        let refgraph = extract_refgraph(&program).unwrap();
+
+        assert_eq!(3, refgraph.declarations.len());
+        assert_eq!(
+            &["foo", "bar", "myfn"],
+            refgraph
+                .declarations
+                .iter()
+                .map(|(_, id)| id.name.as_str())
+                .collect::<Vec<_>>()
+                .as_slice()
+        );
+
+        assert_eq!(1, refgraph.references.len());
+        assert_eq!(
+            &["foo"],
+            refgraph
+                .references
+                .iter()
+                .map(|(_, id)| id.name.as_str())
+                .collect::<Vec<_>>()
+                .as_slice()
+        );
+
+        assert_eq!(1, refgraph.children.len());
+        let myfn = refgraph.children.first().unwrap();
+
+        assert_eq!(myfn.declarations.len(), 1);
+        assert_eq!(
+            &["foo"],
+            myfn.declarations
+                .iter()
+                .map(|(_, id)| id.name.as_str())
+                .collect::<Vec<_>>()
+                .as_slice()
+        );
+
+        assert_eq!(myfn.references.len(), 2);
+        assert_eq!(
+            &["sin", "foo"],
+            myfn.references
+                .iter()
+                .map(|(_, id)| id.name.as_str())
+                .collect::<Vec<_>>()
+                .as_slice()
+        );
+    }
+
+    #[test]
+    fn ast_eq() {
+        const program_str: &str = "\
+foo = 1
+bar = 2
+
+fn myfn() {
+  sin(foo)
+}
+
+baz = myfn()
+";
+
+        let program = kcl!(program_str);
+
+        // if this fails its because the code above is no longer of the
+        // style of the day
+        assert_eq!(program_str, program.recast(&FormatOptions::new(), 0));
+
+        let refgraph = extract_refgraph(&program).unwrap();
+        let edges = refgraph.edges().into_iter().collect::<Vec<_>>();
+
+        // lets say "foo" changed here
+        let tainted_node = program.body.get(0).unwrap();
+
+        let program_min = &program; // minimize_ast(&program, &[tainted_node]);
+        assert_eq!(
+            "\
+foo = 1
+bar = 2
+
+fn myfn() {
+  sin(foo)
+}
+
+baz = myfn()
+",
+            program_min.recast(&FormatOptions::new(), 0)
+        );
+    }
+}

src/wasm-lib/kcl/src/refgraph/walk.rs

@@ -0,0 +1 @@
+

src/wasm-lib/kcl/src/walk/ast_node.rs

@@ -41,6 +41,46 @@ pub enum Node<'a> {
     KclNone(&'a types::KclNone),
 }
 
+impl Node<'_> {
+    ///
+    pub fn ptr_eq(&self, other: Node) -> bool {
+        unsafe { std::ptr::eq(self.ptr(), other.ptr()) }
+    }
+
+    unsafe fn ptr(&self) -> *const () {
+        match self {
+            Node::Program(n) => *n as *const _ as *const (),
+            Node::ImportStatement(n) => *n as *const _ as *const (),
+            Node::ExpressionStatement(n) => *n as *const _ as *const (),
+            Node::VariableDeclaration(n) => *n as *const _ as *const (),
+            Node::ReturnStatement(n) => *n as *const _ as *const (),
+            Node::VariableDeclarator(n) => *n as *const _ as *const (),
+            Node::Literal(n) => *n as *const _ as *const (),
+            Node::TagDeclarator(n) => *n as *const _ as *const (),
+            Node::Identifier(n) => *n as *const _ as *const (),
+            Node::BinaryExpression(n) => *n as *const _ as *const (),
+            Node::FunctionExpression(n) => *n as *const _ as *const (),
+            Node::CallExpression(n) => *n as *const _ as *const (),
+            Node::CallExpressionKw(n) => *n as *const _ as *const (),
+            Node::PipeExpression(n) => *n as *const _ as *const (),
+            Node::PipeSubstitution(n) => *n as *const _ as *const (),
+            Node::ArrayExpression(n) => *n as *const _ as *const (),
+            Node::ArrayRangeExpression(n) => *n as *const _ as *const (),
+            Node::ObjectExpression(n) => *n as *const _ as *const (),
+            Node::MemberExpression(n) => *n as *const _ as *const (),
+            Node::UnaryExpression(n) => *n as *const _ as *const (),
+            Node::Parameter(p) => *p as *const _ as *const (),
+            Node::ObjectProperty(n) => *n as *const _ as *const (),
+            Node::IfExpression(n) => *n as *const _ as *const (),
+            Node::ElseIf(n) => *n as *const _ as *const (),
+            Node::KclNone(n) => *n as *const _ as *const (),
+            Node::LabelledExpression(n) => *n as *const _ as *const (),
+        }
+    }
+
+    // pub fn ptr_eq(&self, other: Node<'_>) -> bool {}
+}
+
 impl Node<'_> {
     /// Return the digest of the [Node], pulling the underlying Digest from
     /// the AST types.

src/wasm-lib/kcl/src/walk/mod.rs

@@ -3,4 +3,5 @@ mod ast_visitor;
 mod ast_walk;
 
 pub use ast_node::Node;
+pub use ast_visitor::{Visitable, Visitor};
 pub use ast_walk::walk;