Skip to main content

spreadsheet_mcp/analysis/
formula.rs

1use crate::model::FormulaGroup;
2use crate::utils::column_number_to_name;
3use anyhow::{Context, Result};
4use formualizer_parse::{
5    ASTNode,
6    parser::{BatchParser, CollectPolicy, ReferenceType},
7    pretty::canonical_formula,
8};
9use parking_lot::{Mutex, RwLock};
10use std::collections::{HashMap, HashSet};
11use std::sync::Arc;
12use umya_spreadsheet::{CellFormulaValues, Worksheet};
13
14const RANGE_EXPANSION_LIMIT: usize = 500;
15
16#[derive(Clone)]
17pub struct FormulaAtlas {
18    parser: Arc<Mutex<BatchParser>>,
19    cache: Arc<RwLock<HashMap<String, Arc<ParsedFormula>>>>,
20    _volatility: Arc<Vec<String>>,
21}
22
23#[derive(Debug, Clone)]
24pub struct ParsedFormula {
25    pub fingerprint: String,
26    pub canonical: String,
27    pub is_volatile: bool,
28    pub dependencies: Vec<String>,
29}
30
31impl FormulaAtlas {
32    pub fn new(volatility_functions: Vec<String>) -> Self {
33        let normalized: Vec<String> = volatility_functions
34            .into_iter()
35            .map(|s| s.to_ascii_uppercase())
36            .collect();
37        let lookup = Arc::new(normalized);
38        let closure_lookup = lookup.clone();
39        let parser = BatchParser::builder()
40            .with_volatility_classifier(move |name| {
41                let upper = name.to_ascii_uppercase();
42                closure_lookup.iter().any(|entry| entry == &upper)
43            })
44            .build();
45        Self {
46            parser: Arc::new(Mutex::new(parser)),
47            cache: Arc::new(RwLock::new(HashMap::new())),
48            _volatility: lookup,
49        }
50    }
51
52    pub fn parse(&self, formula: &str) -> Result<Arc<ParsedFormula>> {
53        if let Some(existing) = self.cache.read().get(formula) {
54            return Ok(existing.clone());
55        }
56
57        let ast = {
58            let mut parser = self.parser.lock();
59            parser
60                .parse(formula)
61                .with_context(|| format!("failed to parse formula: {formula}"))?
62        };
63        let parsed = Arc::new(parsed_from_ast(&ast));
64
65        self.cache
66            .write()
67            .insert(formula.to_string(), parsed.clone());
68        Ok(parsed)
69    }
70}
71
72impl Default for FormulaAtlas {
73    fn default() -> Self {
74        Self::new(default_volatility_functions())
75    }
76}
77
78fn unescape_formula_string(s: &str) -> String {
79    s.replace("\"\"", "\"")
80}
81
82fn parsed_from_ast(ast: &ASTNode) -> ParsedFormula {
83    let fingerprint = format!("{:016x}", ast.fingerprint());
84    let canonical = unescape_formula_string(&canonical_formula(ast));
85    let dependencies = ast
86        .get_dependencies()
87        .iter()
88        .map(|reference| reference_to_string(reference))
89        .collect();
90    ParsedFormula {
91        fingerprint,
92        canonical,
93        is_volatile: ast.contains_volatile(),
94        dependencies,
95    }
96}
97
98pub struct FormulaGraph {
99    precedents: HashMap<String, Vec<String>>,
100    dependents: HashMap<String, Vec<String>>,
101    groups: HashMap<String, FormulaGroupAccumulator>,
102    range_dependents: Vec<RangeDependentEntry>,
103    sheet_name: String,
104}
105
106#[derive(Debug, Clone)]
107struct RangeDependentEntry {
108    #[allow(dead_code)]
109    range_key: String,
110    reference: ReferenceType,
111    dependents: Vec<String>,
112}
113
114impl FormulaGraph {
115    pub fn build(sheet: &Worksheet, atlas: &FormulaAtlas) -> Result<Self> {
116        let sheet_name = sheet.get_name().to_string();
117        let mut precedents_build: HashMap<String, HashSet<String>> = HashMap::new();
118        let mut dependents_build: HashMap<String, HashSet<String>> = HashMap::new();
119        let mut groups: HashMap<String, FormulaGroupAccumulator> = HashMap::new();
120        let mut range_dependents_build: HashMap<String, (ReferenceType, HashSet<String>)> =
121            HashMap::new();
122
123        let collect_policy = CollectPolicy {
124            expand_small_ranges: true,
125            range_expansion_limit: RANGE_EXPANSION_LIMIT,
126            include_names: true,
127        };
128
129        for cell in sheet.get_cell_collection() {
130            if !cell.is_formula() {
131                continue;
132            }
133            let formula_text = cell.get_formula();
134            if formula_text.is_empty() {
135                continue;
136            }
137            let formula_with_prefix = if formula_text.starts_with('=') {
138                formula_text.to_string()
139            } else {
140                format!("={}", formula_text)
141            };
142
143            let ast = {
144                let mut parser = atlas.parser.lock();
145                parser
146                    .parse(&formula_with_prefix)
147                    .with_context(|| format!("failed to parse formula: {formula_with_prefix}"))?
148            };
149
150            let fingerprint = format!("{:016x}", ast.fingerprint());
151            let canonical = unescape_formula_string(&canonical_formula(&ast));
152            let is_volatile = ast.contains_volatile();
153
154            let coordinate = cell.get_coordinate();
155            let address = coordinate.get_coordinate();
156
157            let (is_array, is_shared_type) = cell
158                .get_formula_obj()
159                .map(|obj| match obj.get_formula_type() {
160                    CellFormulaValues::Array => (true, false),
161                    CellFormulaValues::Shared => (false, true),
162                    _ => (false, false),
163                })
164                .unwrap_or((false, false));
165
166            let group =
167                groups
168                    .entry(fingerprint.clone())
169                    .or_insert_with(|| FormulaGroupAccumulator {
170                        canonical: canonical.clone(),
171                        addresses: Vec::new(),
172                        is_volatile,
173                        is_array,
174                        is_shared: is_shared_type,
175                    });
176            if cell.get_formula_shared_index().is_some() {
177                group.is_shared = true;
178            }
179            group.addresses.push(address.clone());
180            group.is_volatile |= is_volatile;
181
182            let refs = ast.collect_references(&collect_policy);
183            for reference in refs {
184                match &reference {
185                    ReferenceType::Cell {
186                        sheet, row, col, ..
187                    } => {
188                        let dep_addr = format_cell_address(sheet.as_deref(), *row, *col);
189                        precedents_build
190                            .entry(address.clone())
191                            .or_default()
192                            .insert(dep_addr.clone());
193                        dependents_build
194                            .entry(dep_addr)
195                            .or_default()
196                            .insert(address.clone());
197                    }
198                    ReferenceType::Range {
199                        start_row,
200                        start_col,
201                        end_row,
202                        end_col,
203                        ..
204                    } => {
205                        let prec_str = reference.to_string();
206                        precedents_build
207                            .entry(address.clone())
208                            .or_default()
209                            .insert(prec_str.clone());
210
211                        if is_large_or_infinite_range(*start_row, *start_col, *end_row, *end_col) {
212                            range_dependents_build
213                                .entry(prec_str)
214                                .or_insert_with(|| (reference.clone(), HashSet::new()))
215                                .1
216                                .insert(address.clone());
217                        }
218                    }
219                    ReferenceType::NamedRange(name) => {
220                        precedents_build
221                            .entry(address.clone())
222                            .or_default()
223                            .insert(name.clone());
224                    }
225                    ReferenceType::Table(_) | ReferenceType::External(_) => {
226                        let table_str = reference.to_string();
227                        precedents_build
228                            .entry(address.clone())
229                            .or_default()
230                            .insert(table_str);
231                    }
232                }
233            }
234        }
235
236        let precedents = precedents_build
237            .into_iter()
238            .map(|(k, v)| (k, v.into_iter().collect()))
239            .collect();
240        let dependents = dependents_build
241            .into_iter()
242            .map(|(k, v)| (k, v.into_iter().collect()))
243            .collect();
244        let range_dependents = range_dependents_build
245            .into_iter()
246            .map(|(key, (ref_type, addrs))| RangeDependentEntry {
247                range_key: key,
248                reference: ref_type,
249                dependents: addrs.into_iter().collect(),
250            })
251            .collect();
252
253        Ok(Self {
254            precedents,
255            dependents,
256            groups,
257            range_dependents,
258            sheet_name,
259        })
260    }
261
262    pub fn groups(&self) -> Vec<FormulaGroup> {
263        self.groups
264            .iter()
265            .map(|(fingerprint, group)| FormulaGroup {
266                fingerprint: fingerprint.clone(),
267                addresses: group.addresses.clone(),
268                count: Some(group.addresses.len() as u32),
269                formula: group.canonical.clone(),
270                is_array: group.is_array,
271                is_shared: group.is_shared,
272                is_volatile: group.is_volatile,
273            })
274            .collect()
275    }
276
277    pub fn precedents(&self, address: &str) -> Vec<String> {
278        self.precedents.get(address).cloned().unwrap_or_default()
279    }
280
281    pub fn dependents(&self, address: &str) -> Vec<String> {
282        self.dependents_limited(address, None).0
283    }
284
285    /// Returns cells that depend on the given address, with optional limit.
286    ///
287    /// Returns (dependents, was_truncated). If limit is Some and exceeded,
288    /// was_truncated is true and only limit dependents are returned.
289    ///
290    /// Performance: O(n) where n = number of large range references in the sheet.
291    /// Early exits when limit reached to keep response times bounded.
292    pub fn dependents_limited(&self, address: &str, limit: Option<usize>) -> (Vec<String>, bool) {
293        let mut result = self.dependents.get(address).cloned().unwrap_or_default();
294        let limit = limit.unwrap_or(usize::MAX);
295
296        if result.len() >= limit {
297            result.truncate(limit);
298            return (result, true);
299        }
300
301        if let Some((row, col)) = parse_cell_address(address) {
302            let (query_sheet, _) = split_sheet_prefix(address);
303            'outer: for entry in &self.range_dependents {
304                if range_contains_cell(&entry.reference, query_sheet, &self.sheet_name, row, col) {
305                    for addr in &entry.dependents {
306                        if !result.contains(addr) {
307                            result.push(addr.clone());
308                            if result.len() >= limit {
309                                break 'outer;
310                            }
311                        }
312                    }
313                }
314            }
315        }
316
317        let truncated = result.len() >= limit;
318        (result, truncated)
319    }
320}
321
322fn format_cell_address(sheet: Option<&str>, row: u32, col: u32) -> String {
323    let col_str = column_number_to_name(col);
324    match sheet {
325        Some(s) => format!("{}!{}{}", s, col_str, row),
326        None => format!("{}{}", col_str, row),
327    }
328}
329
330fn is_large_or_infinite_range(
331    start_row: Option<u32>,
332    start_col: Option<u32>,
333    end_row: Option<u32>,
334    end_col: Option<u32>,
335) -> bool {
336    match (start_row, start_col, end_row, end_col) {
337        (Some(sr), Some(sc), Some(er), Some(ec)) => {
338            let rows = er.saturating_sub(sr) + 1;
339            let cols = ec.saturating_sub(sc) + 1;
340            (rows as usize) * (cols as usize) > RANGE_EXPANSION_LIMIT
341        }
342        _ => true,
343    }
344}
345
346fn range_contains_cell(
347    range: &ReferenceType,
348    query_sheet: Option<&str>,
349    current_sheet: &str,
350    row: u32,
351    col: u32,
352) -> bool {
353    match range {
354        ReferenceType::Range {
355            sheet: range_sheet,
356            start_row,
357            start_col,
358            end_row,
359            end_col,
360            ..
361        } => {
362            let range_sheet_name = range_sheet.as_deref().unwrap_or(current_sheet);
363            let query_sheet_name = query_sheet.unwrap_or(current_sheet);
364            if !range_sheet_name.eq_ignore_ascii_case(query_sheet_name) {
365                return false;
366            }
367            let row_ok = match (start_row, end_row) {
368                (Some(sr), Some(er)) => row >= *sr && row <= *er,
369                (Some(sr), None) => row >= *sr,
370                (None, Some(er)) => row <= *er,
371                (None, None) => true,
372            };
373            let col_ok = match (start_col, end_col) {
374                (Some(sc), Some(ec)) => col >= *sc && col <= *ec,
375                (Some(sc), None) => col >= *sc,
376                (None, Some(ec)) => col <= *ec,
377                (None, None) => true,
378            };
379            row_ok && col_ok
380        }
381        _ => false,
382    }
383}
384
385fn parse_cell_address(address: &str) -> Option<(u32, u32)> {
386    let (_, cell_part) = split_sheet_prefix(address);
387    let cell_part = cell_part.trim_start_matches('$');
388
389    let mut col_str = String::new();
390    let mut row_str = String::new();
391
392    for ch in cell_part.chars() {
393        if ch == '$' {
394            continue;
395        }
396        if ch.is_ascii_alphabetic() && row_str.is_empty() {
397            col_str.push(ch.to_ascii_uppercase());
398        } else if ch.is_ascii_digit() {
399            row_str.push(ch);
400        }
401    }
402
403    if col_str.is_empty() || row_str.is_empty() {
404        return None;
405    }
406
407    let col = column_name_to_number(&col_str)?;
408    let row: u32 = row_str.parse().ok()?;
409    Some((row, col))
410}
411
412fn column_name_to_number(name: &str) -> Option<u32> {
413    let mut result: u32 = 0;
414    for ch in name.chars() {
415        if !ch.is_ascii_alphabetic() {
416            return None;
417        }
418        result = result * 26 + (ch.to_ascii_uppercase() as u32 - 'A' as u32 + 1);
419    }
420    Some(result)
421}
422
423fn split_sheet_prefix(address: &str) -> (Option<&str>, &str) {
424    if let Some(idx) = address.find('!') {
425        let sheet = &address[..idx];
426        let sheet = sheet.trim_start_matches('\'').trim_end_matches('\'');
427        let cell = &address[idx + 1..];
428        (Some(sheet), cell)
429    } else {
430        (None, address)
431    }
432}
433
434struct FormulaGroupAccumulator {
435    canonical: String,
436    addresses: Vec<String>,
437    is_volatile: bool,
438    is_array: bool,
439    is_shared: bool,
440}
441
442fn reference_to_string(reference: &ReferenceType) -> String {
443    reference.to_string()
444}
445
446pub fn normalize_cell_reference(sheet_name: &str, row: u32, col: u32) -> String {
447    format!("{}!{}{}", sheet_name, column_number_to_name(col), row)
448}
449
450fn default_volatility_functions() -> Vec<String> {
451    vec![
452        "NOW",
453        "TODAY",
454        "RAND",
455        "RANDBETWEEN",
456        "OFFSET",
457        "INDIRECT",
458        "INFO",
459        "CELL",
460        "AREAS",
461        "INDEX",
462        "MOD",
463        "ROW",
464        "COLUMN",
465        "ROWS",
466        "COLUMNS",
467        "HYPERLINK",
468    ]
469    .into_iter()
470    .map(|s| s.to_string())
471    .collect()
472}