Add support to patch artemis ast file

src/scripts/artemis/ast/mod.rs

@@ -42,6 +42,15 @@ impl ScriptBuilder for AstScriptBuilder {
     fn script_type(&self) -> &'static ScriptType {
         &ScriptType::Artemis
     }
+
+    fn is_this_format(&self, _filename: &str, buf: &[u8], buf_len: usize) -> Option<u8> {
+        let parser = parser::Parser::new(&buf[..buf_len], Encoding::Utf8);
+        if parser.try_parse_header().is_ok() {
+            Some(15)
+        } else {
+            None
+        }
+    }
 }
 
 #[derive(Debug)]
@@ -207,12 +216,174 @@ impl Script for AstScript {
 
     fn import_messages<'a>(
         &'a self,
-        _messages: Vec<Message>,
+        messages: Vec<Message>,
         mut file: Box<dyn WriteSeek + 'a>,
         encoding: Encoding,
-        _replacement: Option<&'a ReplacementTable>,
+        replacement: Option<&'a ReplacementTable>,
     ) -> Result<()> {
-        let ast = self.ast.clone();
+        let mut ast = self.ast.clone();
+        let root = &mut ast.ast;
+        let mut block_name = root["label"]["top"]["block"]
+            .as_string()
+            .ok_or(anyhow::anyhow!("Missing top block name"))?;
+        let mut block = &mut root[block_name];
+        let mut mess = messages.iter();
+        let mut mes = mess.next();
+        let mut lang = self.lang.as_ref().map(|s| s.to_string());
+        loop {
+            if block[Key("savetitle")].is_array() {
+                let lan = lang.as_ref().map(|s| s.as_str()).unwrap_or("text");
+                let m = match mes {
+                    Some(m) => m,
+                    None => return Err(anyhow::anyhow!("Not enough messages.")),
+                };
+                let mut title = m.message.clone();
+                if let Some(repl) = replacement {
+                    for (k, v) in &repl.map {
+                        title = title.replace(k, v);
+                    }
+                }
+                block[Key("savetitle")][lan].set_string(title);
+                mes = mess.next();
+            }
+            if block["text"].is_array() {
+                let lan = match &lang {
+                    Some(l) => l.as_str(),
+                    None => {
+                        for l in block["text"].kv_keys() {
+                            if l != "vo" {
+                                lang = Some(l.to_string());
+                                break;
+                            }
+                        }
+                        match lang {
+                            Some(ref l) => l.as_str(),
+                            // No text found, continue to next block
+                            None => {
+                                block_name = match block["linknext"].as_string() {
+                                    Some(name) => name,
+                                    None => break,
+                                };
+                                block = &mut root[block_name];
+                                continue;
+                            }
+                        }
+                    }
+                };
+                let origin_names: Vec<_> = {
+                    let mut tex = &block["text"][lan];
+                    if tex.is_null() {
+                        for l in block["text"].kv_keys() {
+                            if l != "vo" {
+                                tex = &block["text"][l];
+                                break;
+                            }
+                        }
+                    }
+                    tex.members().map(|m| m["name"].clone()).collect()
+                };
+                let mut arr = Value::new_array();
+                for name in origin_names {
+                    let m = match mes {
+                        Some(m) => m,
+                        None => return Err(anyhow::anyhow!("Not enough messages.")),
+                    };
+                    let mut text = m.message.clone();
+                    if let Some(repl) = replacement {
+                        for (k, v) in &repl.map {
+                            text = text.replace(k, v);
+                        }
+                    }
+                    if !text.ends_with("\n") {
+                        text.push('\n');
+                    }
+                    let mut v = text::TextParser::new(&text.replace("\n", "<rt2>")).parse()?;
+                    if name.is_array() {
+                        let mut n = match &m.name {
+                            Some(n) => n.to_string(),
+                            None => return Err(anyhow::anyhow!("Message name is missing.")),
+                        };
+                        if let Some(repl) = replacement {
+                            for (k, v) in &repl.map {
+                                n = n.replace(k, v);
+                            }
+                        }
+                        v.insert_member(0, Value::new_kv("name", name));
+                        if v["name"].len() <= 1 {
+                            if v["name"][0] != n {
+                                v["name"].push_member(Value::Str(n));
+                            }
+                        } else {
+                            v["name"].last_member_mut().set_string(n);
+                        }
+                    }
+                    arr.push_member(v);
+                    mes = mess.next();
+                }
+                block["text"][lan] = arr;
+            }
+            if block["select"].is_array() {
+                let lan = match &lang {
+                    Some(l) => l.as_str(),
+                    None => {
+                        for l in block["select"].kv_keys() {
+                            if l != "vo" {
+                                lang = Some(l.to_string());
+                                break;
+                            }
+                        }
+                        match lang {
+                            Some(ref l) => l.as_str(),
+                            // No text found, continue to next block
+                            None => {
+                                block_name = match block["linknext"].as_string() {
+                                    Some(name) => name,
+                                    None => break,
+                                };
+                                block = &mut root[block_name];
+                                continue;
+                            }
+                        }
+                    }
+                };
+                let select_count = {
+                    let mut select = &block["select"][lan];
+                    if select.is_null() {
+                        for l in block["select"].kv_keys() {
+                            if l != "vo" {
+                                select = &block["select"][l];
+                                break;
+                            }
+                        }
+                    }
+                    select.len()
+                };
+                let mut new_select = Value::new_array();
+                for _ in 0..select_count {
+                    let m = match mes {
+                        Some(m) => m,
+                        None => return Err(anyhow::anyhow!("Not enough messages.")),
+                    };
+                    let mut select_text = m.message.clone();
+                    if let Some(repl) = replacement {
+                        for (k, v) in &repl.map {
+                            select_text = select_text.replace(k, v);
+                        }
+                    }
+                    new_select.push_member(Value::Str(select_text));
+                    mes = mess.next();
+                }
+                block["select"][lan] = new_select;
+            }
+            block_name = match block["linknext"].as_string() {
+                Some(name) => name,
+                None => break,
+            };
+            block = &mut root[block_name];
+        }
+        if mes.is_some() || mess.next().is_some() {
+            return Err(anyhow::anyhow!("Not all messages were used."));
+        }
         let mut writer = Vec::new();
         let mut dumper = dump::Dumper::new(&mut writer);
         if self.no_indent {

src/scripts/artemis/ast/parser.rs

@@ -14,7 +14,7 @@ pub struct Parser<'a> {
 }
 
 impl<'a> Parser<'a> {
-    pub fn new<S: AsRef<[u8]>>(str: &'a S, encoding: Encoding) -> Self {
+    pub fn new<S: AsRef<[u8]> + ?Sized>(str: &'a S, encoding: Encoding) -> Self {
         let str = str.as_ref();
         Parser {
             str,
@@ -26,6 +26,14 @@ impl<'a> Parser<'a> {
         }
     }
 
+    pub fn try_parse_header(mut self) -> Result<()> {
+        self.erase_whitespace();
+        self.parse_indent(b"astver")?;
+        self.parse_equal()?;
+        self.parse_f64()?;
+        Ok(())
+    }
+
     pub fn parse(mut self) -> Result<AstFile> {
         self.erase_whitespace();
         self.parse_indent(b"astver")?;
@@ -36,7 +44,7 @@ impl<'a> Parser<'a> {
         if self.is_indent(b"astname") {
             self.parse_indent(b"astname")?;
             self.parse_equal()?;
-            astname = Some(self.parse_str()?.to_string());
+            astname = Some(self.parse_any_str()?.to_string());
             self.erase_whitespace();
         }
         self.parse_indent(b"ast")?;
@@ -141,6 +149,15 @@ impl<'a> Parser<'a> {
         }
     }
 
+    fn parse_any_str(&mut self) -> Result<String> {
+        self.erase_whitespace();
+        match self.peek().ok_or(self.error2("unexpected eof"))? {
+            b'"' => self.parse_str(),
+            b'[' => self.parse_raw_str(),
+            _ => self.error("expected string or raw string"),
+        }
+    }
+
     fn parse_f64(&mut self) -> Result<f64> {
         self.erase_whitespace();
         let start = self.pos;

src/scripts/artemis/ast/text.rs

@@ -1,6 +1,19 @@
 use super::types::*;
 use crate::utils::escape::*;
 use anyhow::Result;
+use unicode_segmentation::UnicodeSegmentation;
+
+fn escape_text(s: &str) -> String {
+    let mut escaped = String::with_capacity(s.len());
+    for c in s.chars() {
+        match c {
+            '&' => escaped.push_str("&"),
+            '<' => escaped.push_str("&lt;"),
+            _ => escaped.push(c),
+        }
+    }
+    escaped
+}
 
 pub struct TextGenerator {
     data: String,
@@ -17,7 +30,7 @@ impl TextGenerator {
         for (i, item) in v.members().enumerate() {
             match item {
                 Value::Str(s) => {
-                    self.data.push_str(s);
+                    self.data.push_str(&escape_text(s));
                 }
                 Value::Float(_) => {
                     return Err(anyhow::anyhow!(
@@ -116,3 +129,227 @@ impl TextGenerator {
         Ok(self.data)
     }
 }
+
+pub struct TextParser<'a> {
+    data: Value,
+    text: Vec<&'a str>,
+    pos: usize,
+    len: usize,
+}
+
+impl<'a> TextParser<'a> {
+    pub fn new(str: &'a str) -> Self {
+        let text: Vec<&'a str> = UnicodeSegmentation::graphemes(str, true).collect();
+        let len = text.len();
+        TextParser {
+            data: Value::new_array(),
+            text,
+            pos: 0,
+            len,
+        }
+    }
+
+    pub fn parse(mut self) -> Result<Value> {
+        while let Some(c) = self.peek() {
+            match c {
+                "<" => {
+                    self.parse_array()?;
+                }
+                _ => {
+                    let mut text = String::new();
+                    self.eat_char();
+                    text.push_str(c);
+                    while let Some(b) = self.peek() {
+                        if b == "<" {
+                            break;
+                        }
+                        text.push_str(b);
+                        self.eat_char();
+                    }
+                    if !text.is_empty() {
+                        self.data.push_member(Value::Str(unescape_xml(&text)));
+                    }
+                }
+            }
+        }
+        Ok(self.data)
+    }
+
+    fn parse_array(&mut self) -> Result<()> {
+        let mut arr = Value::new_array();
+        self.parse_indent("<")?;
+        loop {
+            let c = self.peek().ok_or(self.error2("Unexpected eof"))?;
+            match c {
+                ">" => {
+                    self.eat_char();
+                    break;
+                }
+                "-" | "." | "0" | "1" | "2" | "3" | "4" | "5" | "6" | "7" | "8" | "9" => {
+                    arr.push_member(self.parse_any_number()?);
+                }
+                " " => {
+                    self.eat_char();
+                }
+                _ => {
+                    let key = self.parse_key()?;
+                    let v = if self.is_indent("=") {
+                        self.parse_indent("=")?;
+                        Value::KeyVal((key, Box::new(self.parse_str()?)))
+                    } else {
+                        Value::Str(key)
+                    };
+                    arr.push_member(v);
+                }
+            }
+        }
+        self.data.push_member(arr);
+        Ok(())
+    }
+
+    fn parse_any_number(&mut self) -> Result<Value> {
+        self.erase_whitespace();
+        let mut number = String::new();
+        while let Some(c) = self.peek() {
+            if c == "."
+                || c == "-"
+                || c == "0"
+                || c == "1"
+                || c == "2"
+                || c == "3"
+                || c == "4"
+                || c == "5"
+                || c == "6"
+                || c == "7"
+                || c == "8"
+                || c == "9"
+            {
+                number.push_str(c);
+                self.eat_char();
+            } else {
+                break;
+            }
+        }
+        if number.contains(".") {
+            number
+                .parse()
+                .map(Value::Float)
+                .map_err(|e| self.error2(format!("failed to parse f64: {}", e)))
+        } else {
+            number
+                .parse()
+                .map(Value::Int)
+                .map_err(|e| self.error2(format!("failed to parse i64: {}", e)))
+        }
+    }
+
+    fn parse_key(&mut self) -> Result<String> {
+        self.erase_whitespace();
+        let mut key = String::new();
+        while let Some(c) = self.peek() {
+            if c == "=" || c == " " || c == ">" {
+                break;
+            }
+            key.push_str(c);
+            self.eat_char();
+        }
+        if key.is_empty() {
+            return self.error("Expected key, but found nothing");
+        }
+        Ok(key)
+    }
+
+    fn parse_str(&mut self) -> Result<Value> {
+        self.erase_whitespace();
+        self.parse_indent("\"")?;
+        let mut text = String::new();
+        loop {
+            match self.next().ok_or(self.error2("Unexpected eof"))? {
+                "\"" => {
+                    break;
+                }
+                t => {
+                    text.push_str(t);
+                }
+            }
+        }
+        Ok(Value::Str(unescape_xml(&text)))
+    }
+
+    fn eat_char(&mut self) {
+        if self.pos < self.len {
+            self.pos += 1;
+        }
+    }
+
+    fn erase_whitespace(&mut self) {
+        while let Some(c) = self.peek() {
+            if c == " " {
+                self.eat_char();
+            } else {
+                break;
+            }
+        }
+    }
+
+    fn is_indent(&self, indent: &str) -> bool {
+        let mut pos = self.pos;
+        for ident in indent.graphemes(true) {
+            if pos >= self.len || self.text[pos] != ident {
+                return false;
+            }
+            pos += 1;
+        }
+        true
+    }
+
+    fn parse_indent(&mut self, indent: &str) -> Result<()> {
+        for ident in indent.graphemes(true) {
+            match self.next() {
+                Some(c) => {
+                    if c != ident {
+                        return self.error("Unexpected indent");
+                    }
+                }
+                None => return self.error("Unexpected eof"),
+            }
+        }
+        Ok(())
+    }
+
+    fn next(&mut self) -> Option<&'a str> {
+        if self.pos < self.len {
+            let item = self.text[self.pos];
+            self.pos += 1;
+            Some(item)
+        } else {
+            None
+        }
+    }
+
+    fn peek(&self) -> Option<&'a str> {
+        if self.pos < self.len {
+            Some(self.text[self.pos])
+        } else {
+            None
+        }
+    }
+
+    fn error2<T>(&self, msg: T) -> anyhow::Error
+    where
+        T: std::fmt::Display,
+    {
+        anyhow::anyhow!("Failed to parse at position {}: {}", self.pos, msg)
+    }
+
+    fn error<T, A>(&self, msg: T) -> Result<A>
+    where
+        T: std::fmt::Display,
+    {
+        Err(anyhow::anyhow!(
+            "Failed to parse at position {}: {}",
+            self.pos,
+            msg
+        ))
+    }
+}

src/scripts/artemis/ast/types.rs

@@ -1,4 +1,5 @@
 use std::cmp::{PartialEq, PartialOrd};
+use std::convert::From;
 use std::ops::{Deref, Index, IndexMut};
 
 #[derive(Clone, Debug)]
@@ -11,6 +12,12 @@ pub enum Value {
     Null,
 }
 
+impl From<String> for Value {
+    fn from(s: String) -> Self {
+        Value::Str(s)
+    }
+}
+
 /// Reprsents a key in nested arrays.
 /// For example, in the array `{"save", text="test"}`, the key is `"save"`.
 pub struct Key<'a>(pub &'a str);
@@ -67,6 +74,24 @@ impl Value {
         }
     }
 
+    pub fn find_array_mut(&mut self, key: &str) -> &mut Value {
+        match &self {
+            Value::Array(arr) => {
+                for (i, item) in arr.iter().enumerate() {
+                    if &item[0] == key {
+                        return &mut self[i];
+                    }
+                }
+                self.push_member(Value::Array(vec![Value::Str(key.to_string())]));
+                self.last_member_mut()
+            }
+            _ => {
+                *self = Value::Array(vec![Value::Str(key.to_string())]);
+                self.last_member_mut()
+            }
+        }
+    }
+
     pub fn is_array(&self) -> bool {
         matches!(self, Value::Array(_))
     }
@@ -95,6 +120,50 @@ impl Value {
         }
     }
 
+    pub fn last_member(&self) -> &Value {
+        match self {
+            Value::Array(arr) => arr.last().unwrap_or(&NULL),
+            _ => &NULL,
+        }
+    }
+
+    pub fn last_member_mut(&mut self) -> &mut Value {
+        match self {
+            Value::Array(arr) => {
+                if arr.is_empty() {
+                    arr.push(NULL);
+                }
+                arr.last_mut().unwrap()
+            }
+            _ => {
+                *self = Value::Array(vec![NULL]);
+                self.last_member_mut()
+            }
+        }
+    }
+
+    pub fn len(&self) -> usize {
+        match self {
+            Value::Array(arr) => arr.len(),
+            _ => 0,
+        }
+    }
+
+    pub fn insert_member(&mut self, index: usize, value: Value) {
+        match self {
+            Value::Array(arr) => {
+                if index < arr.len() {
+                    arr.insert(index, value);
+                } else {
+                    arr.push(value);
+                }
+            }
+            _ => {
+                *self = Value::Array(vec![value]);
+            }
+        }
+    }
+
     pub fn members<'a>(&'a self) -> Iter<'a> {
         match self {
             Value::Array(arr) => Iter { iter: arr.iter() },
@@ -111,12 +180,30 @@ impl Value {
         }
     }
 
-    pub fn last_member(&self) -> &Value {
+    pub fn new_array() -> Self {
+        Value::Array(Vec::new())
+    }
+
+    pub fn new_kv<S: Into<String>>(key: S, value: Value) -> Self {
+        Value::KeyVal((key.into(), Box::new(value)))
+    }
+
+    pub fn push_member(&mut self, value: Value) {
         match self {
-            Value::Array(arr) => arr.last().unwrap_or(&NULL),
-            _ => &NULL,
+            Value::Array(arr) => arr.push(value),
+            _ => {
+                *self = Value::Array(vec![value]);
+            }
         }
     }
+
+    pub fn set_str<S: AsRef<str> + ?Sized>(&mut self, value: &S) {
+        *self = Value::Str(value.as_ref().to_string());
+    }
+
+    pub fn set_string<S: Into<String>>(&mut self, value: S) {
+        *self = Value::Str(value.into());
+    }
 }
 
 impl Index<usize> for Value {
@@ -267,6 +354,13 @@ impl<'a, 'b> Index<&'b Key<'a>> for Value {
     }
 }
 
+impl<'a, 'b> IndexMut<&'b Key<'a>> for Value {
+    #[inline(always)]
+    fn index_mut(&mut self, key: &'b Key<'a>) -> &mut Self::Output {
+        self.find_array_mut(&key.0)
+    }
+}
+
 impl<'a> Index<Key<'a>> for Value {
     type Output = Value;
 
@@ -276,6 +370,13 @@ impl<'a> Index<Key<'a>> for Value {
     }
 }
 
+impl<'a> IndexMut<Key<'a>> for Value {
+    #[inline(always)]
+    fn index_mut(&mut self, key: Key<'a>) -> &mut Self::Output {
+        self.find_array_mut(&key.0)
+    }
+}
+
 impl PartialEq<str> for Value {
     fn eq(&self, other: &str) -> bool {
         match self {