feat: added FunctionSymbol

Signed-off-by: shruti2522 <[email protected]> fmt check Signed-off-by: shruti2522 <[email protected]> feat: added FunctionSymbol Signed-off-by: shruti2522 <[email protected]> fmt check Signed-off-by: shruti2522 <[email protected]> feat: added FunctionSymbol Signed-off-by: shruti2522 <[email protected]>
kcl-lang · May 31, 2024 · e10434e · e10434e
1 parent e824bd6
commit e10434e
Showing 1 changed file with 155 additions and 0 deletions.
diff --git a/kclvm/sema/src/core/symbol.rs b/kclvm/sema/src/core/symbol.rs
@@ -67,6 +67,7 @@ pub struct SymbolData {
     pub(crate) exprs: Arena<ExpressionSymbol>,
     pub(crate) comments: Arena<CommentSymbol>,
     pub(crate) decorators: Arena<DecoratorSymbol>,
+    pub(crate) functions: Arena<FunctionSymbol>,
 
     pub(crate) symbols_info: SymbolDB,
 }
@@ -180,6 +181,10 @@ impl SymbolData {
                 .decorators
                 .get(id.get_id())
                 .map(|symbol| symbol as &KCLSymbol),
+            SymbolKind::Function => self
+                .functions
+                .get(id.get_id())
+                .map(|symbol| symbol as &KCLSymbol),
         }
     }
 
@@ -245,6 +250,12 @@ impl SymbolData {
                     symbol
                 });
             }
+            SymbolKind::Function => {
+                self.functions.get_mut(id.get_id()).map(|symbol| {
+                    symbol.sema_info.ty = Some(ty);
+                    symbol
+                });
+            }
         }
     }
 
@@ -690,6 +701,29 @@ impl SymbolData {
         Some(symbol_ref)
     }
 
+    pub fn alloc_function_symbol(
+        &mut self,
+        function: FunctionSymbol,
+        node_key: NodeKey,
+    ) -> SymbolRef {
+        self.symbols_info
+            .symbol_pos_set
+            .insert(function.end.clone());
+        let symbol_id = self.functions.insert(function);
+        let symbol_ref = SymbolRef {
+            id: symbol_id,
+            kind: SymbolKind::Function,
+        };
+        self.symbols_info
+            .node_symbol_map
+            .insert(node_key.clone(), symbol_ref);
+        self.symbols_info
+            .symbol_node_map
+            .insert(symbol_ref, node_key);
+        self.functions.get_mut(symbol_id).unwrap().id = Some(symbol_ref);
+        symbol_ref
+    }
+
     #[inline]
     pub fn get_node_symbol_map(&self) -> &IndexMap<NodeKey, SymbolRef> {
         &self.symbols_info.node_symbol_map
@@ -716,6 +750,7 @@ pub enum SymbolKind {
     Schema,
     Attribute,
     Value,
+    Function,
     Package,
     TypeAlias,
     Unresolved,
@@ -2070,3 +2105,123 @@ impl DecoratorSymbol {
         self.name.clone()
     }
 }
+
+#[derive(Debug, Clone, Default)]
+pub struct FunctionSymbol {
+    pub id: Option<SymbolRef>,
+    pub name: String,
+    pub parameters: Vec<TypeRef>,
+    pub return_type: Option<Arc<Type>>,
+    pub sema_info: KCLSymbolSemanticInfo,
+    pub start: Position,
+    pub end: Position,
+}
+
+impl Symbol for FunctionSymbol {
+    type SymbolData = SymbolData;
+    type SemanticInfo = KCLSymbolSemanticInfo;
+
+    fn get_sema_info(&self) -> &Self::SemanticInfo {
+        &self.sema_info
+    }
+
+    fn is_global(&self) -> bool {
+        true
+    }
+
+    fn get_range(&self) -> Range {
+        (self.start.clone(), self.end.clone())
+    }
+
+    fn get_owner(&self) -> Option<SymbolRef> {
+        None
+    }
+
+    fn get_definition(&self) -> Option<SymbolRef> {
+        self.id.clone()
+    }
+
+    fn get_name(&self) -> String {
+        self.name.clone()
+    }
+
+    fn get_id(&self) -> Option<SymbolRef> {
+        self.id.clone()
+    }
+
+    fn get_attribute(
+        &self,
+        _name: &str,
+        _data: &Self::SymbolData,
+        _module_info: Option<&ModuleInfo>,
+    ) -> Option<SymbolRef> {
+        None
+    }
+
+    fn has_attribute(
+        &self,
+        _name: &str,
+        _data: &Self::SymbolData,
+        _module_info: Option<&ModuleInfo>,
+    ) -> bool {
+        false
+    }
+
+    fn get_all_attributes(
+        &self,
+        _data: &Self::SymbolData,
+        _module_info: Option<&ModuleInfo>,
+    ) -> Vec<SymbolRef> {
+        vec![]
+    }
+
+    fn simple_dump(&self) -> String {
+        let mut output = "{\n".to_string();
+        output.push_str("\"kind\": \"FunctionSymbol\",\n");
+        output.push_str(&format!(
+            "\"range\": \"{}:{}",
+            self.start.filename, self.start.line
+        ));
+        if let Some(start_col) = self.start.column {
+            output.push_str(&format!(":{}", start_col));
+        }
+
+        output.push_str(&format!(" to {}", self.end.line));
+        if let Some(end_col) = self.end.column {
+            output.push_str(&format!(":{}", end_col));
+        }
+        output.push_str(&format!("name :{}", self.get_name()));
+        output.push_str("\"\n}");
+        output
+    }
+
+    fn full_dump(&self, _data: &Self::SymbolData) -> Option<String> {
+        Some(self.simple_dump())
+    }
+}
+
+impl FunctionSymbol {
+    pub fn new(
+        id: Option<SymbolRef>,
+        name: String,
+        parameters: Vec<TypeRef>,
+        return_type: Option<Arc<Type>>,
+        sema_info: KCLSymbolSemanticInfo,
+        start: Position,
+        end: Position,
+    ) -> Self {
+        Self {
+            id,
+            name,
+            parameters,
+            return_type,
+            sema_info,
+            start,
+            end,
+        }
+    }
+
+    pub fn name(&self) -> String {
+        self.name.clone()
+    }
+}