sqruff_lib/core/rules/
context.rs

1use std::any::{Any, TypeId};
2use std::cell::RefCell;
3use std::rc::Rc;
4
5use ahash::AHashMap;
6use sqruff_lib_core::dialects::base::Dialect;
7use sqruff_lib_core::parser::segments::base::{ErasedSegment, Tables};
8use sqruff_lib_core::templaters::base::TemplatedFile;
9
10use crate::core::config::FluffConfig;
11
12#[derive(Debug)]
13pub struct RuleContext<'a> {
14    pub tables: &'a Tables,
15    pub dialect: &'a Dialect,
16    pub templated_file: Option<TemplatedFile>,
17    pub path: Option<String>,
18    pub config: &'a FluffConfig,
19
20    // These change within a file.
21    /// segment: The segment in question
22    pub segment: ErasedSegment,
23    /// parent_stack: A tuple of the path from the root to this segment.
24    pub parent_stack: Vec<ErasedSegment>,
25    /// raw_stack: All the raw segments so far in the file
26    pub raw_stack: Vec<ErasedSegment>,
27    /// memory: Arbitrary storage for the rule
28    pub memory: Rc<RefCell<AHashMap<TypeId, Box<dyn Any>>>>,
29    /// segment_idx: The index of this segment in the parent
30    pub segment_idx: usize,
31}
32
33pub struct Checkpoint {
34    parent_stack: usize,
35    raw_stack: usize,
36}
37
38impl<'a> RuleContext<'a> {
39    pub fn new(
40        tables: &'a Tables,
41        dialect: &'a Dialect,
42        config: &'a FluffConfig,
43        segment: ErasedSegment,
44    ) -> Self {
45        Self {
46            tables,
47            dialect,
48            config,
49            segment,
50            templated_file: <_>::default(),
51            path: <_>::default(),
52            parent_stack: <_>::default(),
53            raw_stack: <_>::default(),
54            memory: Rc::new(RefCell::new(AHashMap::new())),
55            segment_idx: 0,
56        }
57    }
58
59    pub fn checkpoint(&self) -> Checkpoint {
60        Checkpoint {
61            parent_stack: self.parent_stack.len(),
62            raw_stack: self.raw_stack.len(),
63        }
64    }
65
66    pub fn restore(&mut self, checkpoint: Checkpoint) {
67        self.parent_stack.truncate(checkpoint.parent_stack);
68        self.raw_stack.truncate(checkpoint.raw_stack);
69    }
70
71    pub fn try_get<T: Clone + 'static>(&self) -> Option<T> {
72        let id = TypeId::of::<T>();
73
74        let memory = self.memory.borrow();
75        let value = memory.get(&id)?;
76        let value = value.downcast_ref::<T>()?;
77
78        Some(value.clone())
79    }
80
81    pub fn set<T: 'static>(&self, value: T) {
82        let id = TypeId::of::<T>();
83        self.memory.borrow_mut().insert(id, Box::new(value));
84    }
85
86    pub fn siblings_post(&self) -> Vec<ErasedSegment> {
87        if !self.parent_stack.is_empty() {
88            self.parent_stack.last().unwrap().segments()[self.segment_idx + 1..].to_vec()
89        } else {
90            Vec::new()
91        }
92    }
93}