feat: improve lsp completion, implement completion choices for variab…

…les, maps, fts and nfts

components/clarity-lsp/src/common/requests/completion.rs

@@ -3,33 +3,243 @@ use clarity_repl::clarity::{
     functions::{define::DefineFunctions, NativeFunctions},
     variables::NativeVariables,
     vm::types::BlockInfoProperty,
-    ClarityVersion,
+    ClarityVersion, SymbolicExpression,
 };
 use lazy_static::lazy_static;
 use lsp_types::{
     CompletionItem, CompletionItemKind, Documentation, InsertTextFormat, MarkupContent, MarkupKind,
     Position,
 };
 
+lazy_static! {
+    static ref VAR_FUNCTIONS: Vec<String> = vec![
+        NativeFunctions::SetVar.to_string(),
+        NativeFunctions::FetchVar.to_string(),
+    ];
+    static ref MAP_FUNCTIONS: Vec<String> = vec![
+        NativeFunctions::InsertEntry.to_string(),
+        NativeFunctions::FetchEntry.to_string(),
+        NativeFunctions::SetEntry.to_string(),
+        NativeFunctions::DeleteEntry.to_string(),
+    ];
+    static ref FT_FUNCTIONS: Vec<String> = vec![
+        NativeFunctions::GetTokenBalance.to_string(),
+        NativeFunctions::GetTokenSupply.to_string(),
+        NativeFunctions::BurnToken.to_string(),
+        NativeFunctions::MintToken.to_string(),
+        NativeFunctions::TransferToken.to_string(),
+    ];
+    static ref NFT_FUNCTIONS: Vec<String> = vec![
+        NativeFunctions::GetAssetOwner.to_string(),
+        NativeFunctions::BurnAsset.to_string(),
+        NativeFunctions::MintAsset.to_string(),
+        NativeFunctions::TransferAsset.to_string(),
+    ];
+    pub static ref COMPLETION_ITEMS_CLARITY_1: Vec<CompletionItem> =
+        build_default_native_keywords_list(ClarityVersion::Clarity1);
+    pub static ref COMPLETION_ITEMS_CLARITY_2: Vec<CompletionItem> =
+        build_default_native_keywords_list(ClarityVersion::Clarity2);
+}
+
+#[derive(Clone, Debug, Default)]
+pub struct ContractDefinedData {
+    pub vars: Vec<String>,
+    pub maps: Vec<String>,
+    pub fts: Vec<String>,
+    pub nfts: Vec<String>,
+}
+
+pub fn get_contract_defined_data(
+    expressions: Option<&Vec<SymbolicExpression>>,
+) -> Option<ContractDefinedData> {
+    let mut defined_data = ContractDefinedData {
+        ..Default::default()
+    };
+
+    for expression in expressions? {
+        let (define_function, args) = expression.match_list()?.split_first()?;
+        match DefineFunctions::lookup_by_name(define_function.match_atom()?)? {
+            DefineFunctions::PersistedVariable => defined_data
+                .vars
+                .push(args.first()?.match_atom()?.to_string()),
+            DefineFunctions::Map => defined_data
+                .maps
+                .push(args.first()?.match_atom()?.to_string()),
+            DefineFunctions::FungibleToken => defined_data
+                .fts
+                .push(args.first()?.match_atom()?.to_string()),
+            DefineFunctions::NonFungibleToken => defined_data
+                .nfts
+                .push(args.first()?.match_atom()?.to_string()),
+            _ => (),
+        }
+    }
+    Some(defined_data)
+}
+
+#[cfg(test)]
+mod get_contract_defined_data_tests {
+    use clarity_repl::clarity::ast::build_ast_with_rules;
+    use clarity_repl::clarity::stacks_common::types::StacksEpochId;
+    use clarity_repl::clarity::{vm::types::QualifiedContractIdentifier, ClarityVersion};
+
+    use super::{get_contract_defined_data, ContractDefinedData};
+
+    fn get_defined_data(source: &str) -> Option<ContractDefinedData> {
+        let contract_ast = build_ast_with_rules(
+            &QualifiedContractIdentifier::transient(),
+            source,
+            &mut (),
+            ClarityVersion::Clarity2,
+            StacksEpochId::Epoch21,
+            clarity_repl::clarity::ast::ASTRules::Typical,
+        )
+        .unwrap();
+        get_contract_defined_data(Some(&contract_ast.expressions))
+    }
+
+    #[test]
+    fn get_data_vars() {
+        let data = get_defined_data(
+            "(define-data-var counter uint u1) (define-data-var is-active bool true)",
+        )
+        .unwrap_or_default();
+        assert_eq!(data.vars, ["counter", "is-active"]);
+    }
+
+    #[test]
+    fn get_map() {
+        let data = get_defined_data("(define-map names principal { name: (buff 48) })")
+            .unwrap_or_default();
+        assert_eq!(data.maps, ["names"]);
+    }
+
+    #[test]
+    fn get_fts() {
+        let data = get_defined_data("(define-fungible-token clarity-coin)").unwrap_or_default();
+        assert_eq!(data.fts, ["clarity-coin"]);
+    }
+
+    #[test]
+    fn get_nfts() {
+        let data =
+            get_defined_data("(define-non-fungible-token bitcoin-nft uint)").unwrap_or_default();
+        assert_eq!(data.nfts, ["bitcoin-nft"]);
+    }
+}
+
+pub fn populate_snippet_with_options(
+    name: &String,
+    snippet: &String,
+    defined_data: &ContractDefinedData,
+) -> String {
+    if VAR_FUNCTIONS.contains(name) && defined_data.vars.len() > 0 {
+        let choices = defined_data.vars.join(",");
+        return snippet.replace("${1:var}", &format!("${{1|{:}|}}", choices));
+    } else if MAP_FUNCTIONS.contains(name) && defined_data.maps.len() > 0 {
+        let choices = defined_data.maps.join(",");
+        return snippet.replace("${1:map-name}", &format!("${{1|{:}|}}", choices));
+    } else if FT_FUNCTIONS.contains(name) && defined_data.fts.len() > 0 {
+        let choices = defined_data.fts.join(",");
+        return snippet.replace("${1:token-name}", &format!("${{1|{:}|}}", choices));
+    } else if NFT_FUNCTIONS.contains(name) && defined_data.nfts.len() > 0 {
+        let choices = defined_data.nfts.join(",");
+        return snippet.replace("${1:asset-name}", &format!("${{1|{:}|}}", choices));
+    }
+    return snippet.to_string();
+}
+
+#[cfg(test)]
+mod populate_snippet_with_options_tests {
+    use clarity_repl::clarity::ast::build_ast_with_rules;
+    use clarity_repl::clarity::stacks_common::types::StacksEpochId;
+    use clarity_repl::clarity::{vm::types::QualifiedContractIdentifier, ClarityVersion};
+
+    use super::{get_contract_defined_data, populate_snippet_with_options, ContractDefinedData};
+
+    fn get_defined_data(source: &str) -> Option<ContractDefinedData> {
+        let contract_ast = build_ast_with_rules(
+            &QualifiedContractIdentifier::transient(),
+            source,
+            &mut (),
+            ClarityVersion::Clarity2,
+            StacksEpochId::Epoch21,
+            clarity_repl::clarity::ast::ASTRules::Typical,
+        )
+        .unwrap();
+        get_contract_defined_data(Some(&contract_ast.expressions))
+    }
+
+    #[test]
+    fn get_data_vars_snippet() {
+        let data = get_defined_data(
+            "(define-data-var counter uint u1) (define-data-var is-active bool true)",
+        )
+        .unwrap_or_default();
+
+        let snippet = populate_snippet_with_options(
+            &"var-get".to_string(),
+            &"var-get ${1:var}".to_string(),
+            &data,
+        );
+        assert_eq!(snippet, "var-get ${1|counter,is-active|}");
+    }
+
+    #[test]
+    fn get_map_snippet() {
+        let data = get_defined_data("(define-map names principal { name: (buff 48) })")
+            .unwrap_or_default();
+
+        let snippet = populate_snippet_with_options(
+            &"map-get?".to_string(),
+            &"map-get? ${1:map-name} ${2:key-tuple}".to_string(),
+            &data,
+        );
+        assert_eq!(snippet, "map-get? ${1|names|} ${2:key-tuple}");
+    }
+
+    #[test]
+    fn get_fts_snippet() {
+        let data = get_defined_data("(define-fungible-token btc u21)").unwrap_or_default();
+        let snippet = populate_snippet_with_options(
+            &"ft-mint?".to_string(),
+            &"ft-mint? ${1:token-name} ${2:amount} ${3:recipient}".to_string(),
+            &data,
+        );
+        assert_eq!(snippet, "ft-mint? ${1|btc|} ${2:amount} ${3:recipient}");
+    }
+
+    #[test]
+    fn get_nfts_snippet() {
+        let data =
+            get_defined_data("(define-non-fungible-token bitcoin-nft uint)").unwrap_or_default();
+        let snippet = populate_snippet_with_options(
+            &"nft-mint?".to_string(),
+            &"nft-mint? ${1:asset-name} ${2:asset-identifier} ${3:recipient}".to_string(),
+            &data,
+        );
+        assert_eq!(
+            snippet,
+            "nft-mint? ${1|bitcoin-nft|} ${2:asset-identifier} ${3:recipient}"
+        );
+    }
+}
+
 pub fn check_if_should_wrap(source: &str, position: &Position) -> bool {
     if let Some(line) = source
         .lines()
         .collect::<Vec<&str>>()
         .get(position.line as usize)
     {
         let mut chars = line.chars();
-        let mut index = position.character as usize;
-        while index > 0 {
-            index -= 1;
-
-            match chars.nth(index) {
-                Some('(') => return false,
-                Some(char) => {
+        while let Some(char) = chars.next_back() {
+            match char {
+                '(' => return false,
+                char => {
                     if char.is_whitespace() {
                         return true;
                     }
                 }
-                None => return true,
             }
         }
     }
@@ -167,10 +377,3 @@ pub fn build_default_native_keywords_list(version: ClarityVersion) -> Vec<Comple
     .flatten()
     .collect::<Vec<CompletionItem>>()
 }
-
-lazy_static! {
-    pub static ref COMPLETION_ITEMS_CLARITY_1: Vec<CompletionItem> =
-        build_default_native_keywords_list(ClarityVersion::Clarity1);
-    pub static ref COMPLETION_ITEMS_CLARITY_2: Vec<CompletionItem> =
-        build_default_native_keywords_list(ClarityVersion::Clarity2);
-}

components/clarity-lsp/src/common/requests/definitions.rs

@@ -20,7 +20,7 @@ pub enum DefinitionLocation {
 // `global` holds all of the top-level user-defined keywords that are available in the global scope
 // `local` holds the locally user-defined keywords: function parameters, let and match bindings
 // when a user-defined keyword is used in the code, its position and definition location are stored in `tokens`
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Default)]
 pub struct Definitions {
     pub tokens: HashMap<(u32, u32), DefinitionLocation>,
     global: HashMap<ClarityName, Range>,
@@ -31,10 +31,8 @@ pub struct Definitions {
 impl<'a> Definitions {
     pub fn new(deployer: Option<StandardPrincipalData>) -> Self {
         Self {
-            tokens: HashMap::new(),
-            global: HashMap::new(),
-            local: HashMap::new(),
             deployer,
+            ..Default::default()
         }
     }
 

components/clarity-lsp/src/common/state.rs

@@ -27,7 +27,10 @@ use std::collections::{BTreeMap, HashMap, HashSet};
 use std::vec;
 
 use super::requests::capabilities::InitializationOptions;
-use super::requests::completion::{COMPLETION_ITEMS_CLARITY_1, COMPLETION_ITEMS_CLARITY_2};
+use super::requests::completion::{
+    get_contract_defined_data, populate_snippet_with_options, COMPLETION_ITEMS_CLARITY_1,
+    COMPLETION_ITEMS_CLARITY_2,
+};
 use super::requests::definitions::{get_definitions, DefinitionLocation};
 use super::requests::document_symbols::ASTSymbols;
 use super::requests::helpers::{get_atom_start_at_position, get_public_function_definitions};
@@ -317,6 +320,9 @@ impl EditorState {
             .and_then(|p| Some(p.get_completion_items_for_contract(contract_location)))
             .unwrap_or_default();
 
+        let contract_defined_data =
+            get_contract_defined_data(contract.expressions.as_ref()).unwrap_or_default();
+
         let mut completion_items = vec![];
         for mut item in [native_keywords, user_defined_keywords].concat().drain(..) {
             match item.kind {
@@ -326,10 +332,19 @@ impl EditorState {
                     | CompletionItemKind::MODULE
                     | CompletionItemKind::CLASS,
                 ) => {
+                    let mut snippet = item.insert_text.take().unwrap();
+                    if item.kind == Some(CompletionItemKind::FUNCTION) {
+                        snippet = populate_snippet_with_options(
+                            &item.label,
+                            &snippet,
+                            &contract_defined_data,
+                        );
+                    }
+
                     item.insert_text = if should_wrap {
-                        Some(format!("({})", item.insert_text.take().unwrap()))
+                        Some(format!("({})", snippet))
                     } else {
-                        Some(item.insert_text.take().unwrap())
+                        Some(snippet)
                     };
                 }
                 _ => {}

components/clarity-vscode/package.json

@@ -85,7 +85,7 @@
           },
           "clarity-lsp.completionSmartParenthesisWrap": {
             "type": "boolean",
-            "default": false,
+            "default": true,
             "order": 0,
             "description": "If set to true, the completion won't wrap a function in parenthesis if an opening parenthesis is already there"
           },