From 1c9ca1b002badb511ff21d58b3f20986741ab31f Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Sun, 24 Dec 2023 16:04:18 +0800
Subject: feat: custom REPL prompt (#283)

---
 src/utils/mod.rs           |   2 +
 src/utils/render_prompt.rs | 155 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 157 insertions(+)
 create mode 100644 src/utils/render_prompt.rs

(limited to 'src/utils')

diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index bd9f381..04c5668 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -1,11 +1,13 @@
 mod abort_signal;
 mod clipboard;
 mod prompt_input;
+mod render_prompt;
 mod tiktoken;
 
 pub use self::abort_signal::{create_abort_signal, AbortSignal};
 pub use self::clipboard::set_text;
 pub use self::prompt_input::*;
+pub use self::render_prompt::render_prompt;
 pub use self::tiktoken::cl100k_base_singleton;
 
 use sha2::{Digest, Sha256};
diff --git a/src/utils/render_prompt.rs b/src/utils/render_prompt.rs
new file mode 100644
index 0000000..12661fa
--- /dev/null
+++ b/src/utils/render_prompt.rs
@@ -0,0 +1,155 @@
+use std::collections::HashMap;
+
+/// Render REPL prompt
+///
+/// The template comprises plain text and `{...}`.
+///
+/// The syntax of `{...}`:
+/// - `{var}` - When `var` has a value, replace `var` with the value and eval `template`
+/// - `{?var <template>}` - Eval `template` when `var` is evaluated as true
+/// - `{!var <template>}` - Eval `template` when `var` is evaluated as false
+pub fn render_prompt(template: &str, variables: &HashMap<&str, String>) -> String {
+    let exprs = parse_template(template);
+    eval_exprs(&exprs, variables)
+}
+
+fn parse_template(template: &str) -> Vec<Expr> {
+    let chars: Vec<char> = template.chars().collect();
+    let mut exprs = vec![];
+    let mut current = vec![];
+    let mut balances = vec![];
+    for ch in chars.iter().cloned() {
+        if !balances.is_empty() {
+            if ch == '}' {
+                balances.pop();
+                if balances.is_empty() {
+                    if !current.is_empty() {
+                        let block = parse_block(&mut current);
+                        exprs.push(block)
+                    }
+                } else {
+                    current.push(ch);
+                }
+            } else if ch == '{' {
+                balances.push(ch);
+                current.push(ch);
+            } else {
+                current.push(ch);
+            }
+        } else if ch == '{' {
+            balances.push(ch);
+            add_text(&mut exprs, &mut current);
+        } else {
+            current.push(ch)
+        }
+    }
+    add_text(&mut exprs, &mut current);
+    exprs
+}
+
+fn parse_block(current: &mut Vec<char>) -> Expr {
+    let value: String = current.drain(..).collect();
+    match value.split_once(' ') {
+        Some((name, tail)) => {
+            if let Some(name) = name.strip_prefix('?') {
+                let block_exprs = parse_template(tail);
+                Expr::Block(BlockType::Yes, name.to_string(), block_exprs)
+            } else if let Some(name) = name.strip_prefix('!') {
+                let block_exprs = parse_template(tail);
+                Expr::Block(BlockType::No, name.to_string(), block_exprs)
+            } else {
+                Expr::Text(format!("{{{value}}}"))
+            }
+        }
+        None => Expr::Variable(value),
+    }
+}
+
+fn eval_exprs(exprs: &[Expr], variables: &HashMap<&str, String>) -> String {
+    let mut output = String::new();
+    for part in exprs {
+        match part {
+            Expr::Text(text) => output.push_str(text),
+            Expr::Variable(variable) => {
+                let value = variables
+                    .get(variable.as_str())
+                    .cloned()
+                    .unwrap_or_default();
+                output.push_str(&value);
+            }
+            Expr::Block(typ, variable, block_exprs) => {
+                let value = variables
+                    .get(variable.as_str())
+                    .cloned()
+                    .unwrap_or_default();
+                match typ {
+                    BlockType::Yes => {
+                        if truly(&value) {
+                            let block_output = eval_exprs(block_exprs, variables);
+                            output.push_str(&block_output)
+                        }
+                    }
+                    BlockType::No => {
+                        if !truly(&value) {
+                            let block_output = eval_exprs(block_exprs, variables);
+                            output.push_str(&block_output)
+                        }
+                    }
+                }
+            }
+        }
+    }
+    output
+}
+
+fn add_text(exprs: &mut Vec<Expr>, current: &mut Vec<char>) {
+    if current.is_empty() {
+        return;
+    }
+    let value: String = current.drain(..).collect();
+    exprs.push(Expr::Text(value));
+}
+
+fn truly(value: &str) -> bool {
+    !(value.is_empty() || value == "0" || value == "false")
+}
+
+#[derive(Debug)]
+enum Expr {
+    Text(String),
+    Variable(String),
+    Block(BlockType, String, Vec<Expr>),
+}
+
+#[derive(Debug)]
+enum BlockType {
+    Yes,
+    No,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    macro_rules! assert_render {
+        ($template:expr, [$(($key:literal, $value:literal),)*], $expect:literal) => {
+            let data = HashMap::from([
+                $(($key, $value.into()),)*
+            ]);
+            assert_eq!(render_prompt($template, &data), $expect);
+        };
+    }
+
+    #[test]
+    fn test_render() {
+        let prompt = "{?session {session}{?role /}}{role}{?session )}{!session >}";
+        assert_render!(prompt, [], ">");
+        assert_render!(prompt, [("role", "coder"),], "coder>");
+        assert_render!(prompt, [("session", "temp"),], "temp)");
+        assert_render!(
+            prompt,
+            [("session", "temp"), ("role", "coder"),],
+            "temp/coder)"
+        );
+    }
+}
-- 
cgit v1.2.3