Added theme writing code, really bad race condition

crates/settings/src/settings.rs

@@ -12,8 +12,9 @@ use schemars::{
 };
 use serde::{de::DeserializeOwned, Deserialize};
 use serde_json::Value;
-use std::{collections::HashMap, num::NonZeroU32, str, sync::Arc};
+use std::{collections::HashMap, fmt::Write as _, num::NonZeroU32, str, sync::Arc};
 use theme::{Theme, ThemeRegistry};
+use tree_sitter::Query;
 use util::ResultExt as _;
 
 pub use keymap_file::{keymap_file_json_schema, KeymapFileContent};
@@ -501,6 +502,92 @@ pub fn settings_file_json_schema(
     serde_json::to_value(root_schema).unwrap()
 }
 
+pub fn write_theme(mut settings_content: String, new_val: &str) -> String {
+    let mut parser = tree_sitter::Parser::new();
+    parser.set_language(tree_sitter_json::language()).unwrap();
+    let tree = parser.parse(&settings_content, None).unwrap();
+
+    let mut cursor = tree_sitter::QueryCursor::new();
+
+    let query = Query::new(
+        tree_sitter_json::language(),
+        "
+        (document
+            (object
+                (pair
+                    key: (string) @key
+                    value: (_) @value)))
+    ",
+    )
+    .unwrap();
+
+    let mut first_key_start = None;
+    let mut existing_value_range = None;
+    let matches = cursor.matches(&query, tree.root_node(), settings_content.as_bytes());
+    for mat in matches {
+        if mat.captures.len() != 2 {
+            continue;
+        }
+
+        let key = mat.captures[0];
+        let value = mat.captures[1];
+
+        first_key_start.get_or_insert_with(|| key.node.start_byte());
+
+        if let Some(key_text) = settings_content.get(key.node.byte_range()) {
+            if key_text == "\"theme\"" {
+                existing_value_range = Some(value.node.byte_range());
+                break;
+            }
+        }
+    }
+
+    match (first_key_start, existing_value_range) {
+        (None, None) => {
+            // No document, create a new object and overwrite
+            settings_content.clear();
+            write!(settings_content, "{{\n    \"theme\": \"{new_val}\"\n}}\n").unwrap();
+        }
+
+        (_, Some(existing_value_range)) => {
+            // Existing theme key, overwrite
+            settings_content.replace_range(existing_value_range, &format!("\"{new_val}\""));
+        }
+
+        (Some(first_key_start), None) => {
+            // No existing theme key, but other settings. Prepend new theme settings and
+            // match style of first key
+            let mut row = 0;
+            let mut column = 0;
+            for (ix, char) in settings_content.char_indices() {
+                if ix == first_key_start {
+                    break;
+                }
+                if char == '\n' {
+                    row += 1;
+                    column = 0;
+                } else {
+                    column += char.len_utf8();
+                }
+            }
+
+            let content = format!(r#""theme": "{new_val}","#);
+            settings_content.insert_str(first_key_start, &content);
+
+            if row > 0 {
+                settings_content.insert_str(
+                    first_key_start + content.len(),
+                    &format!("\n{:width$}", ' ', width = column),
+                )
+            } else {
+                settings_content.insert_str(first_key_start + content.len(), " ")
+            }
+        }
+    }
+
+    settings_content
+}
+
 fn merge<T: Copy>(target: &mut T, value: Option<T>) {
     if let Some(value) = value {
         *target = value;
@@ -512,3 +599,108 @@ pub fn parse_json_with_comments<T: DeserializeOwned>(content: &str) -> Result<T>
         json_comments::CommentSettings::c_style().strip_comments(content.as_bytes()),
     )?)
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::write_theme;
+    use unindent::Unindent;
+
+    #[test]
+    fn test_write_theme_into_settings_with_theme() {
+        let settings = r#"
+            {
+                "theme": "one-dark"
+            }
+        "#
+        .unindent();
+
+        let new_settings = r#"
+            {
+                "theme": "summerfruit-light"
+            }
+        "#
+        .unindent();
+
+        let settings_after_theme = write_theme(settings, "summerfruit-light");
+
+        assert_eq!(settings_after_theme, new_settings)
+    }
+
+    #[test]
+    fn test_write_theme_into_empty_settings() {
+        let settings = r#"
+            {
+            }
+        "#
+        .unindent();
+
+        let new_settings = r#"
+            {
+                "theme": "summerfruit-light"
+            }
+        "#
+        .unindent();
+
+        let settings_after_theme = write_theme(settings, "summerfruit-light");
+
+        assert_eq!(settings_after_theme, new_settings)
+    }
+
+    #[test]
+    fn test_write_theme_into_no_settings() {
+        let settings = "".to_string();
+
+        let new_settings = r#"
+            {
+                "theme": "summerfruit-light"
+            }
+        "#
+        .unindent();
+
+        let settings_after_theme = write_theme(settings, "summerfruit-light");
+
+        assert_eq!(settings_after_theme, new_settings)
+    }
+
+    #[test]
+    fn test_write_theme_into_single_line_settings_without_theme() {
+        let settings = r#"{ "a": "", "ok": true }"#.to_string();
+        let new_settings = r#"{ "theme": "summerfruit-light", "a": "", "ok": true }"#;
+
+        let settings_after_theme = write_theme(settings, "summerfruit-light");
+
+        assert_eq!(settings_after_theme, new_settings)
+    }
+
+    #[test]
+    fn test_write_theme_pre_object_whitespace() {
+        let settings = r#"          { "a": "", "ok": true }"#.to_string();
+        let new_settings = r#"          { "theme": "summerfruit-light", "a": "", "ok": true }"#;
+
+        let settings_after_theme = write_theme(settings, "summerfruit-light");
+
+        assert_eq!(settings_after_theme, new_settings)
+    }
+
+    #[test]
+    fn test_write_theme_into_multi_line_settings_without_theme() {
+        let settings = r#"
+            {
+                "a": "b"
+            }
+        "#
+        .unindent();
+
+        let new_settings = r#"
+            {
+                "theme": "summerfruit-light",
+                "a": "b"
+            }
+        "#
+        .unindent();
+
+        let settings_after_theme = write_theme(settings, "summerfruit-light");
+
+        assert_eq!(settings_after_theme, new_settings)
+    }
+}