feat: scope api for multiple lang sdks (#996)

* feat: add scope API for multigual SDKs Signed-off-by: peefy <[email protected]> * feat: impl AST index serializer Signed-off-by: peefy <[email protected]> * feat: ignore ast id serialize when flag is false Signed-off-by: peefy <[email protected]> --------- Signed-off-by: peefy <[email protected]>
kcl-lang · Jan 23, 2024 · 38ba669 · 38ba669
1 parent 29fa8d9
commit 38ba669
Show file tree

Hide file tree

Showing 11 changed files with 269 additions and 33 deletions.
diff --git a/kclvm/Cargo.lock b/kclvm/Cargo.lock
diff --git a/kclvm/api/src/service/into.rs b/kclvm/api/src/service/into.rs
@@ -1,10 +1,11 @@
 use crate::gpyrpc::{
-    CliConfig, Error, KeyValuePair, LoadSettingsFilesResult, Message, Position, Symbol, SymbolIndex,
+    CliConfig, Error, KeyValuePair, LoadSettingsFilesResult, Message, Position, Scope, ScopeIndex,
+    Symbol, SymbolIndex,
 };
 use kclvm_config::settings::SettingsFile;
 use kclvm_error::Diagnostic;
-use kclvm_loader::SymbolInfo;
-use kclvm_sema::core::symbol::SymbolRef;
+use kclvm_loader::{ScopeInfo, SymbolInfo};
+use kclvm_sema::core::{scope::ScopeRef, symbol::SymbolRef};
 
 pub(crate) trait IntoLoadSettingsFiles {
     /// Convert self into the LoadSettingsFiles structure.
@@ -23,6 +24,14 @@ pub(crate) trait IntoSymbol {
     fn into_symbol(self) -> Symbol;
 }
 
+pub(crate) trait IntoScope {
+    fn into_scope(self) -> Scope;
+}
+
+pub(crate) trait IntoScopeIndex {
+    fn into_scope_index(self) -> ScopeIndex;
+}
+
 impl IntoLoadSettingsFiles for SettingsFile {
     fn into_load_settings_files(self, files: &[String]) -> LoadSettingsFilesResult {
         LoadSettingsFilesResult {
@@ -89,6 +98,17 @@ impl IntoSymbolIndex for SymbolRef {
     }
 }
 
+impl IntoScopeIndex for ScopeRef {
+    fn into_scope_index(self) -> ScopeIndex {
+        let (index, generation) = self.get_id().into_raw_parts();
+        ScopeIndex {
+            i: index as u64,
+            g: generation as u64,
+            kind: format!("{:?}", self.get_kind()),
+        }
+    }
+}
+
 impl IntoSymbol for SymbolInfo {
     fn into_symbol(self) -> Symbol {
         Symbol {
@@ -101,3 +121,15 @@ impl IntoSymbol for SymbolInfo {
         }
     }
 }
+
+impl IntoScope for ScopeInfo {
+    fn into_scope(self) -> Scope {
+        Scope {
+            kind: format!("{:?}", self.kind),
+            parent: self.parent.map(|o| o.into_scope_index()),
+            owner: self.owner.map(|o| o.into_symbol_index()),
+            children: self.children.iter().map(|a| a.into_scope_index()).collect(),
+            defs: self.defs.iter().map(|a| a.into_symbol_index()).collect(),
+        }
+    }
+}
diff --git a/kclvm/api/src/service/service_impl.rs b/kclvm/api/src/service/service_impl.rs
@@ -173,6 +173,8 @@ impl KclvmServiceImpl {
     /// assert_eq!(result.type_errors.len(), 0);
     /// assert_eq!(result.node_symbol_map.len(), 159);
     /// assert_eq!(result.symbols.len(), 12);
+    /// assert_eq!(result.scopes.len(), 3);
+    /// assert_eq!(result.pkg_scope_map.len(), 3);
     /// ```
     pub fn load_package(&self, args: &LoadPackageArgs) -> anyhow::Result<LoadPackageResult> {
         let mut package_maps = HashMap::new();
@@ -191,16 +193,29 @@ impl KclvmServiceImpl {
             resolve_ast: args.resolve_ast,
             load_builtin: args.load_builtin,
         })?;
+        if args.with_ast_index {
+            // Thread local options
+            kclvm_ast::ast::set_should_serialize_id(true);
+        }
         let program_json = serde_json::to_string(&packages.program)?;
         let mut node_symbol_map = HashMap::new();
+        let mut pkg_scope_map = HashMap::new();
         let mut symbols = HashMap::new();
+        let mut scopes = HashMap::new();
         for (k, s) in packages.node_symbol_map {
             node_symbol_map.insert(k.id.to_string(), s.into_symbol_index());
         }
+        for (k, s) in packages.pkg_scope_map {
+            pkg_scope_map.insert(k, s.into_scope_index());
+        }
         for (k, s) in packages.symbols {
             let symbol_index_string = serde_json::to_string(&k)?;
             symbols.insert(symbol_index_string, s.into_symbol());
         }
+        for (k, s) in packages.scopes {
+            let scope_index_string = serde_json::to_string(&k)?;
+            scopes.insert(scope_index_string, s.into_scope());
+        }
         Ok(LoadPackageResult {
             program: program_json,
             paths: packages
@@ -209,7 +224,9 @@ impl KclvmServiceImpl {
                 .map(|p| p.to_str().unwrap().to_string())
                 .collect(),
             node_symbol_map,
+            pkg_scope_map,
             symbols,
+            scopes,
             parse_errors: packages
                 .parse_errors
                 .into_iter()

diff --git a/kclvm/ast/Cargo.toml b/kclvm/ast/Cargo.toml
@@ -13,6 +13,7 @@ serde_json = "1.0"
 
 kclvm-span = { path = "../span" }
 kclvm-error = { path = "../error" }
+thread_local = "1.1.7"
 
 [dev-dependencies]
 kclvm-parser = { path = "../parser" }
diff --git a/kclvm/ast/src/ast.rs b/kclvm/ast/src/ast.rs
@@ -33,7 +33,7 @@
 //! in the compiler and regenerate the walker code.
 //! :copyright: Copyright The KCL Authors. All rights reserved.
 
-use serde::{Deserialize, Serialize};
+use serde::{ser::SerializeStruct, Deserialize, Serialize, Serializer};
 use std::collections::HashMap;
 
 use compiler_base_span::{Loc, Span};
@@ -43,6 +43,11 @@ use uuid;
 use super::token;
 use crate::{node_ref, pos::ContainsPos};
 use kclvm_error::{diagnostic::Range, Position};
+use std::cell::RefCell;
+
+thread_local! {
+    static SHOULD_SERIALIZE_ID: RefCell<bool> = RefCell::new(false);
+}
 
 /// PosTuple denotes the tuple `(filename, line, column, end_line, end_column)`.
 pub type PosTuple = (String, u64, u64, u64, u64);
@@ -87,7 +92,7 @@ impl Serialize for AstIndex {
     where
         S: serde::Serializer,
     {
-        serializer.serialize_bytes(self.0.as_bytes())
+        serializer.serialize_str(&self.to_string())
     }
 }
 
@@ -107,9 +112,9 @@ impl ToString for AstIndex {
 /// that all AST nodes need to contain.
 /// In fact, column and end_column are the counts of character,
 /// For example, `\t` is counted as 1 character, so it is recorded as 1 here, but generally col is 4.
-#[derive(Serialize, Deserialize, Clone, PartialEq)]
+#[derive(Deserialize, Clone, PartialEq)]
 pub struct Node<T> {
-    #[serde(skip_serializing, skip_deserializing, default)]
+    #[serde(serialize_with = "serialize_id", skip_deserializing, default)]
     pub id: AstIndex,
     pub node: T,
     pub filename: String,
@@ -119,6 +124,33 @@ pub struct Node<T> {
     pub end_column: u64,
 }
 
+impl<T: Serialize> Serialize for Node<T> {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        let should_serialize_id = SHOULD_SERIALIZE_ID.with(|f| *f.borrow());
+        let mut state =
+            serializer.serialize_struct("Node", if should_serialize_id { 7 } else { 6 })?;
+        if should_serialize_id {
+            state.serialize_field("id", &self.id)?;
+        }
+        state.serialize_field("node", &self.node)?;
+        state.serialize_field("filename", &self.filename)?;
+        state.serialize_field("line", &self.line)?;
+        state.serialize_field("column", &self.column)?;
+        state.serialize_field("end_line", &self.end_line)?;
+        state.serialize_field("end_column", &self.end_column)?;
+        state.end()
+    }
+}
+
+pub fn set_should_serialize_id(value: bool) {
+    SHOULD_SERIALIZE_ID.with(|f| {
+        *f.borrow_mut() = value;
+    });
+}
+
 impl<T: Debug> Debug for Node<T> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("Node")

diff --git a/kclvm/loader/src/lib.rs b/kclvm/loader/src/lib.rs
@@ -7,7 +7,11 @@ use kclvm_error::{diagnostic::Range, Diagnostic};
 use kclvm_parser::{load_program, KCLModuleCache, LoadProgramOptions, ParseSessionRef};
 use kclvm_sema::{
     advanced_resolver::AdvancedResolver,
-    core::{global_state::GlobalState, symbol::SymbolRef},
+    core::{
+        global_state::GlobalState,
+        scope::{LocalSymbolScopeKind, ScopeData, ScopeRef},
+        symbol::{SymbolData, SymbolRef},
+    },
     namer::Namer,
     resolver::{resolve_program_with_opts, scope::NodeKey},
     ty::{Type, TypeRef},
@@ -42,12 +46,16 @@ pub struct Packages {
     pub paths: Vec<PathBuf>,
     /// All Parse errors
     pub parse_errors: Errors,
-    // Type errors
+    /// Type errors
     pub type_errors: Errors,
-    // Symbol-Type mapping
+    /// Symbol information
     pub symbols: IndexMap<SymbolRef, SymbolInfo>,
+    /// Scope information
+    pub scopes: IndexMap<ScopeRef, ScopeInfo>,
     // AST Node-Symbol mapping
     pub node_symbol_map: IndexMap<NodeKey, SymbolRef>,
+    // <Package path>-<Root scope> mapping
+    pub pkg_scope_map: IndexMap<String, ScopeRef>,
 }
 
 #[derive(Debug, Clone)]
@@ -61,6 +69,33 @@ pub struct SymbolInfo {
     pub is_global: bool,
 }
 
+#[derive(Debug, Clone)]
+pub struct ScopeInfo {
+    /// Scope kind
+    pub kind: ScopeKind,
+    /// Scope parent
+    pub parent: Option<ScopeRef>,
+    /// Scope owner
+    pub owner: Option<SymbolRef>,
+    /// Children scopes
+    pub children: Vec<ScopeRef>,
+    /// Definitions
+    pub defs: Vec<SymbolRef>,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum ScopeKind {
+    Package,
+    Module,
+    List,
+    Dict,
+    Quant,
+    Lambda,
+    SchemaDef,
+    SchemaConfig,
+    Value,
+}
+
 /// load_package provides users with the ability to parse kcl program and sematic model
 /// information including symbols, types, definitions, etc.
 pub fn load_packages(opts: &LoadPackageOptions) -> Result<Packages> {
@@ -101,8 +136,7 @@ pub fn load_packages(opts: &LoadPackageOptions) -> Result<Packages> {
         paths: parse_result.paths,
         parse_errors,
         type_errors,
-        symbols: IndexMap::new(),
-        node_symbol_map: IndexMap::new(),
+        ..Default::default()
     };
     if !opts.resolve_ast {
         return Ok(packages);
@@ -143,5 +177,66 @@ pub fn load_packages(opts: &LoadPackageOptions) -> Result<Packages> {
             }
         }
     }
+    let scopes = gs.get_scopes();
+    for (path, scope_ref) in scopes.get_root_scope_map() {
+        packages.pkg_scope_map.insert(path.clone(), *scope_ref);
+        // Root scopes
+        if let Some(scope_ref) = scopes.get_root_scope(path.clone()) {
+            collect_scope_info(
+                &mut packages.scopes,
+                &scope_ref,
+                scopes,
+                symbols,
+                ScopeKind::Package,
+            );
+        }
+    }
     Ok(packages)
 }
+
+impl From<LocalSymbolScopeKind> for ScopeKind {
+    fn from(value: LocalSymbolScopeKind) -> Self {
+        match value {
+            LocalSymbolScopeKind::List => ScopeKind::List,
+            LocalSymbolScopeKind::Dict => ScopeKind::Dict,
+            LocalSymbolScopeKind::Quant => ScopeKind::Quant,
+            LocalSymbolScopeKind::Lambda => ScopeKind::Lambda,
+            LocalSymbolScopeKind::SchemaDef => ScopeKind::SchemaDef,
+            LocalSymbolScopeKind::SchemaConfig => ScopeKind::SchemaConfig,
+            LocalSymbolScopeKind::Value => ScopeKind::Value,
+        }
+    }
+}
+
+fn collect_scope_info(
+    scopes: &mut IndexMap<ScopeRef, ScopeInfo>,
+    scope_ref: &ScopeRef,
+    scope_data: &ScopeData,
+    symbol_data: &SymbolData,
+    kind: ScopeKind,
+) {
+    if let Some(scope) = scope_data.get_scope(scope_ref) {
+        let kind = if let Some(scope) = scope_data.try_get_local_scope(scope_ref) {
+            scope.get_kind().clone().into()
+        } else {
+            kind
+        };
+        scopes.insert(
+            scope_ref.clone(),
+            ScopeInfo {
+                kind,
+                parent: scope.get_parent(),
+                owner: scope.get_owner(),
+                children: scope.get_children(),
+                defs: scope
+                    .get_all_defs(scope_data, symbol_data, None, false)
+                    .values()
+                    .copied()
+                    .collect::<Vec<_>>(),
+            },
+        );
+        for s in scope.get_children() {
+            collect_scope_info(scopes, &s, scope_data, symbol_data, ScopeKind::Module);
+        }
+    }
+}
diff --git a/kclvm/sema/src/advanced_resolver/mod.rs b/kclvm/sema/src/advanced_resolver/mod.rs
@@ -75,9 +75,9 @@ pub struct Context<'ctx> {
     cur_node: AstIndex,
 
     // whether the identifier currently being visited may be a definition
-    // it will only be true when visiting a lvalue or parameter,
+    // it will only be true when visiting a l-value or parameter,
     // which means advanced resolver will will create the corresponding
-    // ValueSymbol instead of an UnresolveSymbol
+    // ValueSymbol instead of an UnresolvedSymbol
     maybe_def: bool,
 }
 

diff --git a/kclvm/sema/src/advanced_resolver/node.rs b/kclvm/sema/src/advanced_resolver/node.rs
@@ -413,7 +413,7 @@ impl<'ctx> MutSelfTypedResultWalker<'ctx> for AdvancedResolver<'ctx> {
             .node_symbol_map
             .get(&self.ctx.get_node_key(&schema_attr.name.id))?;
         let parent_scope = *self.ctx.scopes.last().unwrap();
-        let parent_scope = self.gs.get_scopes().get_scope(parent_scope).unwrap();
+        let parent_scope = self.gs.get_scopes().get_scope(&parent_scope).unwrap();
         let mut doc = None;
         if let Some(schema_symbol) = parent_scope.get_owner() {
             let schema_symbol = self.gs.get_symbols().get_symbol(schema_symbol).unwrap();