feat(grit): support for Grit pattern, predicate and function definiti…

…ons (#3983)

crates/biome_grit_patterns/src/errors.rs

@@ -18,9 +18,18 @@ pub enum CompileError {
     /// A metavariables was discovered in an unexpected context.
     UnexpectedMetavariable,
 
+    /// If a function with the same name is defined multiple times.
+    DuplicateFunctionDefinition(String),
+
     /// If a function or bubble pattern has multiple parameters with the same name.
     DuplicateParameters,
 
+    /// If a function with the same name is defined multiple times.
+    DuplicatePatternDefinition(String),
+
+    /// If a function with the same name is defined multiple times.
+    DuplicatePredicateDefinition(String),
+
     /// A metavariable was expected at the given range.
     InvalidMetavariableRange(ByteRange),
 
@@ -89,9 +98,18 @@ impl Diagnostic for CompileError {
             CompileError::UnexpectedMetavariable => {
                 fmt.write_markup(markup! { "Unexpected metavariable" })
             }
+            CompileError::DuplicateFunctionDefinition(name) => {
+                fmt.write_markup(markup! { "Duplicate function definition: "{{name}} })
+            }
             CompileError::DuplicateParameters => {
                 fmt.write_markup(markup! { "Duplicate parameters" })
             }
+            CompileError::DuplicatePatternDefinition(name) => {
+                fmt.write_markup(markup! { "Duplicate pattern definition: "{{name}} })
+            }
+            CompileError::DuplicatePredicateDefinition(name) => {
+                fmt.write_markup(markup! { "Duplicate predicate definition: "{{name}} })
+            }
             CompileError::InvalidMetavariableRange(_) => {
                 fmt.write_markup(markup! { "Invalid range for metavariable" })
             }

crates/biome_grit_patterns/src/grit_context.rs

@@ -45,9 +45,9 @@ pub struct GritExecContext<'a> {
 
     loadable_files: &'a [GritTargetFile],
     files: &'a FileOwners<GritTargetTree>,
-    functions: Vec<GritFunctionDefinition<GritQueryContext>>,
-    patterns: Vec<PatternDefinition<GritQueryContext>>,
-    predicates: Vec<PredicateDefinition<GritQueryContext>>,
+    functions: &'a [GritFunctionDefinition<GritQueryContext>],
+    patterns: &'a [PatternDefinition<GritQueryContext>],
+    predicates: &'a [PredicateDefinition<GritQueryContext>],
 }
 
 impl<'a> GritExecContext<'a> {
@@ -56,30 +56,33 @@ impl<'a> GritExecContext<'a> {
         name: Option<&'a str>,
         loadable_files: &'a [GritTargetFile],
         files: &'a FileOwners<GritTargetTree>,
+        functions: &'a [GritFunctionDefinition<GritQueryContext>],
+        patterns: &'a [PatternDefinition<GritQueryContext>],
+        predicates: &'a [PredicateDefinition<GritQueryContext>],
     ) -> Self {
         Self {
             lang,
             name,
             loadable_files,
             files,
-            functions: Vec::new(),
-            patterns: Vec::new(),
-            predicates: Vec::new(),
+            functions,
+            patterns,
+            predicates,
         }
     }
 }
 
 impl<'a> ExecContext<'a, GritQueryContext> for GritExecContext<'a> {
     fn pattern_definitions(&self) -> &[PatternDefinition<GritQueryContext>] {
-        &self.patterns
+        self.patterns
     }
 
     fn predicate_definitions(&self) -> &[PredicateDefinition<GritQueryContext>] {
-        &self.predicates
+        self.predicates
     }
 
     fn function_definitions(&self) -> &[GritFunctionDefinition<GritQueryContext>] {
-        &self.functions
+        self.functions
     }
 
     fn ignore_limit_pattern(&self) -> bool {

crates/biome_grit_patterns/src/grit_definitions.rs

@@ -0,0 +1,167 @@
+use std::collections::BTreeMap;
+
+use crate::{
+    grit_context::GritQueryContext,
+    pattern_compiler::{
+        compilation_context::{DefinitionInfo, NodeCompilationContext},
+        FunctionDefinitionCompiler, PatternDefinitionCompiler, PredicateDefinitionCompiler,
+    },
+    util::TextRangeGritExt,
+    CompileError,
+};
+use biome_grit_syntax::{AnyGritDefinition, GritDefinitionList, GritVariableList};
+use biome_rowan::AstNode;
+use grit_pattern_matcher::pattern::{
+    GritFunctionDefinition, PatternDefinition, PredicateDefinition,
+};
+use grit_util::ByteRange;
+
+#[derive(Clone, Debug)]
+pub struct Definitions {
+    pub patterns: Vec<PatternDefinition<GritQueryContext>>,
+    pub predicates: Vec<PredicateDefinition<GritQueryContext>>,
+    pub functions: Vec<GritFunctionDefinition<GritQueryContext>>,
+}
+
+/// Compiles all definitions.
+///
+/// Must be called after [scan_definitions()].
+pub fn compile_definitions(
+    definitions: GritDefinitionList,
+    context: &mut NodeCompilationContext,
+) -> Result<Definitions, CompileError> {
+    let mut patterns = Vec::new();
+    let mut predicates = Vec::new();
+    let mut functions = Vec::new();
+    for definition in definitions {
+        match definition? {
+            AnyGritDefinition::AnyGritPattern(_) => continue, // Handled separately.
+            AnyGritDefinition::GritPatternDefinition(node) => {
+                patterns.push(PatternDefinitionCompiler::from_node(node, context)?);
+            }
+            AnyGritDefinition::GritPredicateDefinition(node) => {
+                predicates.push(PredicateDefinitionCompiler::from_node(node, context)?);
+            }
+            AnyGritDefinition::GritFunctionDefinition(node) => {
+                functions.push(FunctionDefinitionCompiler::from_node(node, context)?);
+            }
+            AnyGritDefinition::GritBogusDefinition(_) => {
+                unreachable!(); // Should be handled in `scan_definitions()`.
+            }
+        }
+    }
+
+    Ok(Definitions {
+        patterns,
+        predicates,
+        functions,
+    })
+}
+
+pub struct ScannedDefinitionInfo {
+    pub pattern_definition_info: BTreeMap<String, DefinitionInfo>,
+    pub predicate_definition_info: BTreeMap<String, DefinitionInfo>,
+    pub function_definition_info: BTreeMap<String, DefinitionInfo>,
+}
+
+/// Finds all definitions so that we can allocate their scopes in preparation
+/// for the compilation phase.
+pub fn scan_definitions(
+    definitions: GritDefinitionList,
+) -> Result<ScannedDefinitionInfo, CompileError> {
+    let mut pattern_definition_info = BTreeMap::new();
+    let mut pattern_index = 0;
+
+    let mut predicate_definition_info = BTreeMap::new();
+    let mut predicate_index = 0;
+
+    let mut function_definition_info = BTreeMap::new();
+    let mut function_index = 0;
+
+    for definition in definitions {
+        match definition? {
+            AnyGritDefinition::AnyGritPattern(_) => continue, // Handled separately.
+            AnyGritDefinition::GritPatternDefinition(node) => {
+                let name = node.name()?.text();
+                let name = name.trim();
+                if pattern_definition_info.contains_key(name) {
+                    return Err(CompileError::DuplicatePatternDefinition(name.to_owned()));
+                }
+
+                pattern_definition_info.insert(
+                    name.to_owned(),
+                    DefinitionInfo {
+                        index: pattern_index,
+                        parameters: collect_variables(node.args()?.grit_variable_list())?,
+                    },
+                );
+
+                pattern_index += 1;
+            }
+            AnyGritDefinition::GritPredicateDefinition(node) => {
+                let name = node.name()?.text();
+                let name = name.trim();
+                if predicate_definition_info.contains_key(name) {
+                    return Err(CompileError::DuplicatePredicateDefinition(name.to_owned()));
+                }
+
+                predicate_definition_info.insert(
+                    name.to_owned(),
+                    DefinitionInfo {
+                        index: predicate_index,
+                        parameters: collect_variables(node.args()?.grit_variable_list())?,
+                    },
+                );
+
+                predicate_index += 1;
+            }
+            AnyGritDefinition::GritFunctionDefinition(node) => {
+                let name = node.name()?.text();
+                let name = name.trim();
+                if function_definition_info.contains_key(name) {
+                    return Err(CompileError::DuplicateFunctionDefinition(name.to_owned()));
+                }
+
+                function_definition_info.insert(
+                    name.to_owned(),
+                    DefinitionInfo {
+                        index: function_index,
+                        parameters: collect_variables(node.args())?,
+                    },
+                );
+
+                function_index += 1;
+            }
+            AnyGritDefinition::GritBogusDefinition(bogus) => {
+                return Err(CompileError::UnexpectedKind(
+                    bogus
+                        .items()
+                        .next()
+                        .map(|item| item.kind().into())
+                        .unwrap_or_default(),
+                ));
+            }
+        }
+    }
+
+    Ok(ScannedDefinitionInfo {
+        pattern_definition_info,
+        predicate_definition_info,
+        function_definition_info,
+    })
+}
+
+fn collect_variables(
+    variables: GritVariableList,
+) -> Result<Vec<(String, ByteRange)>, CompileError> {
+    variables
+        .into_iter()
+        .map(|var| {
+            let token = var?.value_token()?;
+            Ok((
+                token.text_trimmed().to_string(),
+                token.text_trimmed_range().to_byte_range(),
+            ))
+        })
+        .collect()
+}

crates/biome_grit_patterns/src/grit_query.rs

@@ -1,5 +1,8 @@
 use crate::diagnostics::CompilerDiagnostic;
 use crate::grit_context::{GritExecContext, GritQueryContext, GritTargetFile};
+use crate::grit_definitions::{
+    compile_definitions, scan_definitions, Definitions, ScannedDefinitionInfo,
+};
 use crate::grit_resolved_pattern::GritResolvedPattern;
 use crate::grit_target_language::GritTargetLanguage;
 use crate::grit_tree::GritTargetTree;
@@ -40,6 +43,9 @@ const GLOBAL_VARS: [(&str, usize); 4] = [
 pub struct GritQuery {
     pub pattern: Pattern<GritQueryContext>,
 
+    /// Definitions for named patterns, predicates and functions.
+    pub definitions: Definitions,
+
     /// Diagnostics discovered during compilation of the query.
     pub diagnostics: Vec<CompilerDiagnostic>,
 
@@ -63,6 +69,9 @@ impl GritQuery {
             self.name.as_deref(),
             &files,
             &file_owners,
+            &self.definitions.functions,
+            &self.definitions.patterns,
+            &self.definitions.predicates,
         );
 
         let var_registry = VarRegistry::from_locations(&self.variable_locations);
@@ -91,16 +100,28 @@ impl GritQuery {
 
     pub fn from_node(
         root: GritRoot,
-        path: Option<&Path>,
+        source_path: Option<&Path>,
         lang: GritTargetLanguage,
     ) -> Result<Self, CompileError> {
-        let context = CompilationContext::new(path, lang);
+        let ScannedDefinitionInfo {
+            pattern_definition_info,
+            predicate_definition_info,
+            function_definition_info,
+        } = scan_definitions(root.definitions())?;
+
+        let context = CompilationContext {
+            source_path,
+            lang,
+            pattern_definition_info,
+            predicate_definition_info,
+            function_definition_info,
+        };
 
         let mut vars_array = vec![GLOBAL_VARS
             .iter()
             .map(|global_var| VariableSourceLocations {
                 name: global_var.0.to_string(),
-                file: path
+                file: source_path
                     .map(Path::to_string_lossy)
                     .map_or_else(|| "unnamed".to_owned(), |p| p.to_string()),
                 locations: BTreeSet::new(),
@@ -124,22 +145,23 @@ impl GritQuery {
             &mut diagnostics,
         );
 
+        let mut definitions = compile_definitions(root.definitions(), &mut node_context)?;
+
         let pattern = PatternCompiler::from_node(
             &root.pattern().ok_or(CompileError::MissingPattern)?,
             &mut node_context,
         )?;
 
-        let mut pattern_definitions = Vec::new();
         let pattern = auto_wrap_pattern(
             pattern,
-            &mut pattern_definitions,
+            &mut definitions.patterns,
             true,
             None,
             &mut node_context,
             None,
         )?;
 
-        let name = path
+        let name = source_path
             .and_then(Path::file_stem)
             .map(OsStr::to_string_lossy)
             .map(|stem| stem.into_owned());
@@ -148,6 +170,7 @@ impl GritQuery {
 
         Ok(Self {
             pattern,
+            definitions,
             name,
             language,
             diagnostics,

crates/biome_grit_patterns/src/lib.rs

@@ -4,6 +4,7 @@ mod grit_analysis_ext;
 mod grit_binding;
 mod grit_code_snippet;
 mod grit_context;
+mod grit_definitions;
 mod grit_file;
 mod grit_js_parser;
 mod grit_node;

crates/biome_grit_patterns/src/pattern_compiler.rs

@@ -37,6 +37,7 @@ mod contains_compiler;
 mod divide_compiler;
 mod equal_compiler;
 mod every_compiler;
+mod function_definition_compiler;
 mod if_compiler;
 mod includes_compiler;
 mod like_compiler;
@@ -53,8 +54,10 @@ mod multiply_compiler;
 mod node_like_compiler;
 mod not_compiler;
 mod or_compiler;
+mod pattern_definition_compiler;
 mod predicate_call_compiler;
 mod predicate_compiler;
+mod predicate_definition_compiler;
 mod predicate_return_compiler;
 mod regex_compiler;
 mod rewrite_compiler;
@@ -67,6 +70,10 @@ mod variable_compiler;
 mod where_compiler;
 mod within_compiler;
 
+pub use function_definition_compiler::FunctionDefinitionCompiler;
+pub use pattern_definition_compiler::PatternDefinitionCompiler;
+pub use predicate_definition_compiler::PredicateDefinitionCompiler;
+
 use self::{
     accumulate_compiler::AccumulateCompiler, add_compiler::AddCompiler,
     after_compiler::AfterCompiler, and_compiler::AndCompiler, any_compiler::AnyCompiler,