From 4141c0763e713f9989f18339b43e81da42e1878f Mon Sep 17 00:00:00 2001
From: Morgante Pell <morgante@grit.io>
Date: Fri, 19 Jul 2024 23:52:14 -0700
Subject: [PATCH] feat: add rewritten_text function (#429)

---
 crates/core/src/built_in_functions.rs | 29 ++++++++++++++++-----
 crates/core/src/test.rs               | 37 +++++++++++++++++++++++++++
 2 files changed, 60 insertions(+), 6 deletions(-)

diff --git a/crates/core/src/built_in_functions.rs b/crates/core/src/built_in_functions.rs
index 981a6c07f..96ba3a6f9 100644
--- a/crates/core/src/built_in_functions.rs
+++ b/crates/core/src/built_in_functions.rs
@@ -12,12 +12,12 @@ use grit_pattern_matcher::{
         ResolvedSnippet, State,
     },
 };
-use grit_util::{AnalysisLogs, Language};
+use grit_util::{AnalysisLogs, CodeRange, Language};
 use im::Vector;
 use itertools::Itertools;
 use rand::prelude::SliceRandom;
 use rand::Rng;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashMap};
 
 // todo we can probably use a macro to generate a function that takes a vec and
 // and calls the input function with the vec args unpacked.
@@ -144,7 +144,7 @@ impl BuiltIns {
             BuiltInFunction::new("capitalize", vec!["string"], Box::new(capitalize_fn)),
             BuiltInFunction::new("lowercase", vec!["string"], Box::new(lowercase_fn)),
             BuiltInFunction::new("uppercase", vec!["string"], Box::new(uppercase_fn)),
-            BuiltInFunction::new("text", vec!["string"], Box::new(text_fn)),
+            BuiltInFunction::new("text", vec!["string", "linearize"], Box::new(text_fn)),
             BuiltInFunction::new("trim", vec!["string", "trim_chars"], Box::new(trim_fn)),
             BuiltInFunction::new("join", vec!["list", "separator"], Box::new(join_fn)),
             BuiltInFunction::new("distinct", vec!["list"], Box::new(distinct_fn)),
@@ -245,10 +245,27 @@ fn text_fn<'a>(
 ) -> Result<MarzanoResolvedPattern<'a>> {
     let args = MarzanoResolvedPattern::from_patterns(args, state, context, logs)?;
 
-    let s = match args.first() {
-        Some(Some(resolved_pattern)) => resolved_pattern.text(&state.files, context.language())?,
-        _ => return Err(anyhow!("text takes 1 argument")),
+    let Some(Some(resolved_pattern)) = args.first() else {
+        return Err(anyhow!("text takes 1 argument"));
     };
+    let should_linearize = match args.get(1) {
+        Some(Some(resolved_pattern)) => resolved_pattern.is_truthy(state, context.language())?,
+        _ => false,
+    };
+    if !should_linearize {
+        let text = resolved_pattern.text(&state.files, context.language())?;
+        return Ok(ResolvedPattern::from_string(text.to_string()));
+    }
+    let mut memo: HashMap<CodeRange, Option<String>> = HashMap::new();
+    let effects: Vec<_> = state.effects.clone().into_iter().collect();
+    let s = resolved_pattern.linearized_text(
+        context.language(),
+        &effects,
+        &state.files,
+        &mut memo,
+        false,
+        logs,
+    )?;
     Ok(ResolvedPattern::from_string(s.to_string()))
 }
 
diff --git a/crates/core/src/test.rs b/crates/core/src/test.rs
index 4e29ca5e9..9d7454f4e 100644
--- a/crates/core/src/test.rs
+++ b/crates/core/src/test.rs
@@ -6540,6 +6540,43 @@ fn even_more_linearized_test() {
     .unwrap();
 }
 
+#[test]
+fn linearized_text_fn() {
+    run_test_expected({
+        TestArgExpected {
+            pattern: r#"
+                |language js
+                |
+                |function go_to_the_zoo($a) js {
+                |    return $a.text.replaceAll("baz", "bars").replaceAll("zoo", "alexandria")
+                |}
+                |
+                |program() where {
+                |    $program <: contains bubble `foo($a)` => `zoo($a)`,
+                |    $current_text = text($program, true),
+                |    $final_text = go_to_the_zoo($current_text),
+                |    $program => $final_text
+                |}
+                |"#
+            .trim_margin()
+            .unwrap(),
+            source: r#"
+                |foo(bar(baz))
+                |foo(blo, fly)
+                |"#
+            .trim_margin()
+            .unwrap(),
+            expected: r#"
+                |alexandria(bar(bars))
+                |alexandria(blo, fly)
+                |"#
+            .trim_margin()
+            .unwrap(),
+        }
+    })
+    .unwrap();
+}
+
 #[test]
 fn matching_var_snippet_work() {
     run_test_expected({