Skip to main content

spreadsheet_mcp/formula/
pattern.rs

1use anyhow::{Result, anyhow, bail};
2use formualizer_parse::parser::ReferenceType;
3use formualizer_parse::pretty::canonical_formula;
4use formualizer_parse::{ASTNode, ASTNodeType};
5
6#[derive(Debug, Clone, Copy, PartialEq, Eq)]
7pub enum RelativeMode {
8    Excel,
9    AbsCols,
10    AbsRows,
11}
12
13impl RelativeMode {
14    pub fn parse(mode: Option<&str>) -> Result<Self> {
15        match mode.unwrap_or("excel").to_ascii_lowercase().as_str() {
16            "excel" => Ok(Self::Excel),
17            "abs_cols" | "abscols" | "columns_absolute" => Ok(Self::AbsCols),
18            "abs_rows" | "absrows" | "rows_absolute" => Ok(Self::AbsRows),
19            other => bail!("invalid relative_mode: {}", other),
20        }
21    }
22}
23
24pub fn parse_base_formula(formula: &str) -> Result<ASTNode> {
25    let trimmed = formula.trim();
26    let with_equals = if trimmed.starts_with('=') {
27        trimmed.to_string()
28    } else {
29        format!("={}", trimmed)
30    };
31    formualizer_parse::parse(&with_equals)
32        .map_err(|e| anyhow!("failed to parse base_formula: {}", e.message))
33}
34
35pub fn shift_formula_ast(
36    ast: &ASTNode,
37    delta_col: i32,
38    delta_row: i32,
39    mode: RelativeMode,
40) -> Result<String> {
41    let mut shifted = ast.clone();
42    shift_refs_in_place(&mut shifted, delta_col, delta_row, mode)?;
43    Ok(canonical_formula(&shifted))
44}
45
46/// Walk the AST and mutate all reference nodes in-place.
47fn shift_refs_in_place(
48    node: &mut ASTNode,
49    delta_col: i32,
50    delta_row: i32,
51    mode: RelativeMode,
52) -> Result<()> {
53    match &mut node.node_type {
54        ASTNodeType::Reference {
55            original,
56            reference,
57        } => {
58            shift_reference_in_place(original, reference, delta_col, delta_row, mode)?;
59        }
60        ASTNodeType::UnaryOp { expr, .. } => {
61            shift_refs_in_place(expr, delta_col, delta_row, mode)?;
62        }
63        ASTNodeType::BinaryOp { left, right, .. } => {
64            shift_refs_in_place(left, delta_col, delta_row, mode)?;
65            shift_refs_in_place(right, delta_col, delta_row, mode)?;
66        }
67        ASTNodeType::Function { args, .. } => {
68            for arg in args.iter_mut() {
69                shift_refs_in_place(arg, delta_col, delta_row, mode)?;
70            }
71        }
72        ASTNodeType::Array(rows) => {
73            for row in rows.iter_mut() {
74                for cell in row.iter_mut() {
75                    shift_refs_in_place(cell, delta_col, delta_row, mode)?;
76                }
77            }
78        }
79        ASTNodeType::Literal(_) => {}
80    }
81    Ok(())
82}
83
84fn shift_reference_in_place(
85    original: &mut String,
86    reference: &mut ReferenceType,
87    delta_col: i32,
88    delta_row: i32,
89    mode: RelativeMode,
90) -> Result<()> {
91    match reference {
92        ReferenceType::Cell {
93            row,
94            col,
95            row_abs,
96            col_abs,
97            ..
98        } => {
99            if mode == RelativeMode::AbsCols {
100                *col_abs = true;
101            }
102            if mode == RelativeMode::AbsRows {
103                *row_abs = true;
104            }
105            *col = shift_u32(*col, *col_abs, delta_col)?;
106            *row = shift_u32(*row, *row_abs, delta_row)?;
107        }
108        ReferenceType::Range {
109            start_row,
110            start_col,
111            end_row,
112            end_col,
113            start_row_abs,
114            start_col_abs,
115            end_row_abs,
116            end_col_abs,
117            ..
118        } => {
119            if mode == RelativeMode::AbsCols {
120                if start_col.is_some() {
121                    *start_col_abs = true;
122                }
123                if end_col.is_some() {
124                    *end_col_abs = true;
125                }
126            }
127            if mode == RelativeMode::AbsRows {
128                if start_row.is_some() {
129                    *start_row_abs = true;
130                }
131                if end_row.is_some() {
132                    *end_row_abs = true;
133                }
134            }
135            *start_col = shift_opt_u32(*start_col, *start_col_abs, delta_col)?;
136            *end_col = shift_opt_u32(*end_col, *end_col_abs, delta_col)?;
137            *start_row = shift_opt_u32(*start_row, *start_row_abs, delta_row)?;
138            *end_row = shift_opt_u32(*end_row, *end_row_abs, delta_row)?;
139        }
140        // Table refs and named ranges don't shift
141        ReferenceType::Table(_) | ReferenceType::NamedRange(_) | ReferenceType::External(_) => {}
142    }
143    // Update the original string to match the mutated reference
144    *original = reference.to_string();
145    Ok(())
146}
147
148fn shift_u32(value: u32, abs: bool, delta: i32) -> Result<u32> {
149    if abs || delta == 0 {
150        return Ok(value);
151    }
152    let shifted = value as i64 + delta as i64;
153    if shifted < 1 {
154        bail!("shift would move reference before A1");
155    }
156    Ok(shifted as u32)
157}
158
159fn shift_opt_u32(value: Option<u32>, abs: bool, delta: i32) -> Result<Option<u32>> {
160    match value {
161        Some(v) => Ok(Some(shift_u32(v, abs, delta)?)),
162        None => Ok(None),
163    }
164}