Nested codegen support (#33)

README.md

@@ -37,6 +37,7 @@ asset_dir = "test/"
 write_dir = "output/"
 typescript = true
 luau = true
+style = "flat"
 
 [creator]
 type = "user"
@@ -61,6 +62,8 @@ id = 9670971
     -   Generate a Typescript definition file.
 -   `luau`: boolean (optional)
     -   Use the `luau` file extension.
+-   `style`: string (optional)
+    -   The code-generation style to use. Defaults to `flat`.
 -   `output_name`: string (optional)
     -   The name for the generated files. Defaults to `assets`.
 -   `existing`: map<string, ExistingAsset> (optional)

src/codegen.rs → src/codegen/flat.rs

@@ -47,46 +47,3 @@ pub fn generate_ts(
         output_dir, interface, output_dir
     ))
 }
-
-#[cfg(test)]
-mod tests {
-    use std::collections::BTreeMap;
-
-    use crate::{FileEntry, LockFile};
-
-    fn test_lockfile() -> LockFile {
-        let mut entries = BTreeMap::new();
-        entries.insert(
-            "assets/foo.png".to_string(),
-            FileEntry {
-                asset_id: 1,
-                hash: "a".to_string(),
-            },
-        );
-        entries.insert(
-            "assets/bar/baz.png".to_string(),
-            FileEntry {
-                asset_id: 2,
-                hash: "b".to_string(),
-            },
-        );
-
-        LockFile { entries }
-    }
-
-    #[test]
-    fn generate_lua() {
-        let lockfile = test_lockfile();
-
-        let lua = super::generate_lua(&lockfile, "assets").unwrap();
-        assert_eq!(lua, "return {\n\t[\"/bar/baz.png\"] = \"rbxassetid://2\",\n\t[\"/foo.png\"] = \"rbxassetid://1\"\n}");
-    }
-
-    #[test]
-    fn generate_ts() {
-        let lockfile = test_lockfile();
-
-        let ts = super::generate_ts(&lockfile, "assets", "assets").unwrap();
-        assert_eq!(ts, "declare const assets: {\n\t\"/bar/baz.png\": string,\n\t\"/foo.png\": string\n}\nexport = assets");
-    }
-}

src/codegen/mod.rs

@@ -0,0 +1,89 @@
+use crate::{config::StyleType, LockFile};
+
+mod flat;
+mod nested;
+
+pub fn generate_lua(
+    lockfile: &LockFile,
+    strip_dir: &str,
+    style: &StyleType,
+) -> anyhow::Result<String> {
+    match style {
+        StyleType::Flat => flat::generate_lua(lockfile, strip_dir),
+        StyleType::Nested => nested::generate_lua(lockfile, strip_dir),
+    }
+}
+
+pub fn generate_ts(
+    lockfile: &LockFile,
+    strip_dir: &str,
+    output_dir: &str,
+    style: &StyleType,
+) -> anyhow::Result<String> {
+    match style {
+        StyleType::Flat => flat::generate_ts(lockfile, strip_dir, output_dir),
+        StyleType::Nested => nested::generate_ts(lockfile, strip_dir, output_dir),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::BTreeMap;
+
+    use crate::{FileEntry, LockFile};
+
+    fn test_lockfile() -> LockFile {
+        let mut entries = BTreeMap::new();
+        entries.insert(
+            "assets/foo.png".to_string(),
+            FileEntry {
+                asset_id: 1,
+                hash: "a".to_string(),
+            },
+        );
+        entries.insert(
+            "assets/bar/baz.png".to_string(),
+            FileEntry {
+                asset_id: 2,
+                hash: "b".to_string(),
+            },
+        );
+        LockFile { entries }
+    }
+
+    #[test]
+    fn generate_lua() {
+        let lockfile = test_lockfile();
+
+        let lua = super::flat::generate_lua(&lockfile, "assets").unwrap();
+        assert_eq!(lua, "return {\n\t[\"/bar/baz.png\"] = \"rbxassetid://2\",\n\t[\"/foo.png\"] = \"rbxassetid://1\"\n}");
+    }
+
+    #[test]
+    fn generate_ts() {
+        let lockfile = test_lockfile();
+
+        let ts = super::flat::generate_ts(&lockfile, "assets", "assets").unwrap();
+        assert_eq!(ts, "declare const assets: {\n\t\"/bar/baz.png\": string,\n\t\"/foo.png\": string\n}\nexport = assets");
+    }
+
+    #[test]
+    fn generate_lua_nested() {
+        let lockfile = test_lockfile();
+
+        let lua = super::nested::generate_lua(&lockfile, "assets").unwrap();
+        assert_eq!(
+            lua,
+            "return {\n    bar = {\n        [\"baz.png\"] = \"rbxassetid://2\",\n    },\n    [\"foo.png\"] = \"rbxassetid://1\",\n}");
+    }
+
+    #[test]
+    fn generate_ts_nested() {
+        let lockfile = test_lockfile();
+
+        let ts = super::nested::generate_ts(&lockfile, "assets", "assets").unwrap();
+        assert_eq!(
+            ts,
+            "declare const assets: {\n    bar: {\n        \"baz.png\": \"rbxassetid://2\",\n    },\n    \"foo.png\": \"rbxassetid://1\",\n}\nexport = assets");
+    }
+}

src/codegen/nested/ast.rs

@@ -0,0 +1,233 @@
+use std::fmt::{self, Write};
+
+macro_rules! proxy_display {
+    ( $target: ty ) => {
+        impl fmt::Display for $target {
+            fn fmt(&self, output: &mut fmt::Formatter) -> fmt::Result {
+                let mut stream = AstStream::new(output, &self.1);
+                AstFormat::fmt_ast(self, &mut stream)
+            }
+        }
+    };
+}
+
+trait AstFormat {
+    fn fmt_ast(&self, output: &mut AstStream) -> fmt::Result;
+    fn fmt_key(&self, output: &mut AstStream<'_, '_>) -> fmt::Result {
+        write!(output, "[")?;
+        self.fmt_ast(output)?;
+        write!(output, "]")
+    }
+}
+
+#[derive(Debug)]
+pub(crate) enum AstTarget {
+    Lua,
+    Typescript { output_dir: String },
+}
+
+pub(crate) struct AstStream<'a, 'b> {
+    number_of_spaces: usize,
+    indents: usize,
+    is_start_of_line: bool,
+    writer: &'a mut (dyn Write),
+    target: &'b AstTarget,
+}
+
+impl<'a, 'b> AstStream<'a, 'b> {
+    pub fn new(writer: &'a mut (dyn fmt::Write + 'a), target: &'b AstTarget) -> Self {
+        Self {
+            number_of_spaces: 4,
+            indents: 0,
+            is_start_of_line: true,
+            writer,
+            target,
+        }
+    }
+
+    fn indent(&mut self) {
+        self.indents += 1
+    }
+
+    fn unindent(&mut self) {
+        if self.indents > 0 {
+            self.indents -= 1
+        }
+    }
+
+    fn begin_line(&mut self) -> fmt::Result {
+        self.is_start_of_line = true;
+        self.writer.write_char('\n')
+    }
+}
+
+impl Write for AstStream<'_, '_> {
+    fn write_str(&mut self, value: &str) -> fmt::Result {
+        let mut is_first_line = true;
+
+        for line in value.split('\n') {
+            if is_first_line {
+                is_first_line = false;
+            } else {
+                self.begin_line()?;
+            }
+
+            if !line.is_empty() {
+                if self.is_start_of_line {
+                    self.is_start_of_line = false;
+                    self.writer.write_str(&format!(
+                        "{: >1$}",
+                        "",
+                        self.number_of_spaces * self.indents
+                    ))?;
+                }
+
+                self.writer.write_str(line)?;
+            }
+        }
+
+        Ok(())
+    }
+}
+
+proxy_display!(ReturnStatement);
+
+#[derive(Debug)]
+pub(crate) struct ReturnStatement(pub Expression, pub AstTarget);
+
+impl AstFormat for ReturnStatement {
+    fn fmt_ast(&self, output: &mut AstStream) -> fmt::Result {
+        match output.target {
+            AstTarget::Lua => {
+                write!(output, "return ")
+            }
+            AstTarget::Typescript { output_dir } => {
+                write!(output, "declare const {output_dir}: ")
+            }
+        }?;
+        let result = self.0.fmt_ast(output);
+        if let AstTarget::Typescript { output_dir } = output.target {
+            write!(output, "\nexport = {output_dir}")?
+        }
+        result
+    }
+}
+
+#[derive(Debug)]
+pub(crate) enum Expression {
+    String(String),
+    Table(Table),
+}
+
+impl Expression {
+    pub fn table(expressions: Vec<(Expression, Expression)>) -> Self {
+        Self::Table(Table { expressions })
+    }
+}
+
+impl AstFormat for Expression {
+    fn fmt_ast(&self, output: &mut AstStream) -> fmt::Result {
+        match self {
+            Self::Table(val) => val.fmt_ast(output),
+            Self::String(val) => val.fmt_ast(output),
+        }
+    }
+
+    fn fmt_key(&self, output: &mut AstStream<'_, '_>) -> fmt::Result {
+        match self {
+            Self::Table(val) => val.fmt_key(output),
+            Self::String(val) => val.fmt_key(output),
+        }
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct Table {
+    pub expressions: Vec<(Expression, Expression)>,
+}
+
+impl AstFormat for Table {
+    fn fmt_ast(&self, output: &mut AstStream<'_, '_>) -> fmt::Result {
+        let assignment = match output.target {
+            AstTarget::Lua => " = ",
+            AstTarget::Typescript { .. } => ": ",
+        };
+
+        writeln!(output, "{{")?;
+        output.indent();
+
+        for (key, value) in &self.expressions {
+            key.fmt_key(output)?;
+            write!(output, "{assignment}")?;
+            value.fmt_ast(output)?;
+            writeln!(output, ",")?;
+        }
+
+        output.unindent();
+        write!(output, "}}")
+    }
+}
+
+impl AstFormat for String {
+    fn fmt_ast(&self, output: &mut AstStream) -> fmt::Result {
+        write!(output, "\"{}\"", self)
+    }
+
+    fn fmt_key(&self, output: &mut AstStream<'_, '_>) -> fmt::Result {
+        if is_valid_identifier(self) {
+            write!(output, "{}", self)
+        } else {
+            match output.target {
+                AstTarget::Lua => write!(output, "[\"{}\"]", self),
+                AstTarget::Typescript { .. } => write!(output, "\"{}\"", self),
+            }
+        }
+    }
+}
+
+impl From<String> for Expression {
+    fn from(value: String) -> Self {
+        Self::String(value)
+    }
+}
+
+impl From<&'_ String> for Expression {
+    fn from(value: &String) -> Self {
+        Self::String(value.clone())
+    }
+}
+
+impl From<&'_ str> for Expression {
+    fn from(value: &str) -> Self {
+        Self::String(value.to_owned())
+    }
+}
+
+impl From<Table> for Expression {
+    fn from(value: Table) -> Self {
+        Self::Table(value)
+    }
+}
+
+fn is_valid_ident_char_start(value: char) -> bool {
+    value.is_ascii_alphabetic() || value == '_'
+}
+
+fn is_valid_ident_char(value: char) -> bool {
+    value.is_ascii_alphanumeric() || value == '_'
+}
+
+fn is_valid_identifier(value: &str) -> bool {
+    let mut chars = value.chars();
+
+    match chars.next() {
+        Some(first) => {
+            if !is_valid_ident_char_start(first) {
+                return false;
+            }
+        }
+        None => return false,
+    }
+
+    chars.all(is_valid_ident_char)
+}