Fix Lua snapshot code to work with children.

It's also way easier to read now

rojo-test/build-tests/init_with_children/init.lua

@@ -0,0 +1 @@
+-- init.lua

rojo-test/build-tests/init_with_children/other.lua

@@ -0,0 +1 @@
+-- other.lua

rojo-test/src/build_test.rs

@@ -24,6 +24,7 @@ gen_build_tests! {
     csv_in_folder,
     deep_nesting,
     gitkeep,
+    init_with_children,
     json_model_in_folder,
     json_model_legacy_name,
     module_in_folder,

rojo-test/src/snapshots/build_test__init_with_children.snap

@@ -0,0 +1,20 @@
+---
+created: "2019-08-30T00:12:56.686707700Z"
+creator: insta@0.10.1
+source: rojo-test/src/build_test.rs
+expression: contents
+---
+<roblox version="4">
+  <Item class="ModuleScript" referent="0">
+    <Properties>
+      <string name="Name">init_with_children</string>
+      <string name="Source">-- init.lua</string>
+    </Properties>
+    <Item class="ModuleScript" referent="1">
+      <Properties>
+        <string name="Name">other</string>
+        <string name="Source">-- other.lua</string>
+      </Properties>
+    </Item>
+  </Item>
+</roblox>

src/snapshot_middleware/lua.rs

@@ -8,7 +8,10 @@ use crate::{
     snapshot::InstanceSnapshot,
 };
 
-use super::middleware::{SnapshotFileResult, SnapshotInstanceResult, SnapshotMiddleware};
+use super::{
+    dir::SnapshotDir,
+    middleware::{SnapshotFileResult, SnapshotInstanceResult, SnapshotMiddleware},
+};
 
 pub struct SnapshotLua;
 
@@ -19,64 +22,29 @@ impl SnapshotMiddleware for SnapshotLua {
     ) -> SnapshotInstanceResult<'static> {
         let file_name = entry.path().file_name().unwrap().to_string_lossy();
 
-        if entry.is_directory() {
-            let module_init_path = entry.path().join("init.lua");
-            if let Some(init_entry) = imfs.get(module_init_path).with_not_found()? {
-                if let Some(mut snapshot) = SnapshotLua::from_imfs(imfs, &init_entry)? {
-                    snapshot.name = Cow::Owned(file_name.into_owned());
-
-                    return Ok(Some(snapshot));
-                }
-            }
-
-            let server_init_path = entry.path().join("init.server.lua");
-            if let Some(init_entry) = imfs.get(server_init_path).with_not_found()? {
-                if let Some(mut snapshot) = SnapshotLua::from_imfs(imfs, &init_entry)? {
-                    snapshot.name = Cow::Owned(file_name.into_owned());
-
-                    return Ok(Some(snapshot));
-                }
-            }
-
-            let client_init_path = entry.path().join("init.client.lua");
-            if let Some(init_entry) = imfs.get(client_init_path).with_not_found()? {
-                if let Some(mut snapshot) = SnapshotLua::from_imfs(imfs, &init_entry)? {
-                    snapshot.name = Cow::Owned(file_name.into_owned());
-
-                    return Ok(Some(snapshot));
-                }
-            }
+        // These paths alter their parent instance, so we don't need to turn
+        // them into a script instance here.
+        match &*file_name {
+            "init.lua" | "init.server.lua" | "init.client.lua" => return Ok(None),
+            _ => {}
         }
 
-        let (class_name, instance_name) =
-            if let Some(name) = match_trailing(&file_name, ".server.lua") {
-                ("Script", name)
-            } else if let Some(name) = match_trailing(&file_name, ".client.lua") {
-                ("LocalScript", name)
-            } else if let Some(name) = match_trailing(&file_name, ".lua") {
-                ("ModuleScript", name)
+        if entry.is_file() {
+            snapshot_lua_file(imfs, entry)
+        } else {
+            if let Some(snapshot) = snapshot_init(imfs, entry, "init.lua")? {
+                // An `init.lua` file turns its parent into a ModuleScript
+                Ok(Some(snapshot))
+            } else if let Some(snapshot) = snapshot_init(imfs, entry, "init.server.lua")? {
+                // An `init.server.lua` file turns its parent into a Script
+                Ok(Some(snapshot))
+            } else if let Some(snapshot) = snapshot_init(imfs, entry, "init.client.lua")? {
+                // An `init.client.lua` file turns its parent into a LocalScript
+                Ok(Some(snapshot))
             } else {
-                return Ok(None);
-            };
-
-        let contents = entry.contents(imfs)?;
-        let contents_str = str::from_utf8(contents)
-            .expect("File content was not valid UTF-8")
-            .to_string();
-
-        let properties = hashmap! {
-            "Source".to_owned() => RbxValue::String {
-                value: contents_str,
-            },
-        };
-
-        Ok(Some(InstanceSnapshot {
-            snapshot_id: None,
-            name: Cow::Owned(instance_name.to_owned()),
-            class_name: Cow::Borrowed(class_name),
-            properties,
-            children: Vec::new(),
-        }))
+                Ok(None)
+            }
+        }
     }
 
     fn from_instance(tree: &RbxTree, id: RbxId) -> SnapshotFileResult {
@@ -91,6 +59,70 @@ impl SnapshotMiddleware for SnapshotLua {
     }
 }
 
+/// Core routine for turning Lua files into snapshots.
+fn snapshot_lua_file<F: ImfsFetcher>(
+    imfs: &mut Imfs<F>,
+    entry: &ImfsEntry,
+) -> SnapshotInstanceResult<'static> {
+    let file_name = entry.path().file_name().unwrap().to_string_lossy();
+
+    let (class_name, instance_name) = if let Some(name) = match_trailing(&file_name, ".server.lua")
+    {
+        ("Script", name)
+    } else if let Some(name) = match_trailing(&file_name, ".client.lua") {
+        ("LocalScript", name)
+    } else if let Some(name) = match_trailing(&file_name, ".lua") {
+        ("ModuleScript", name)
+    } else {
+        return Ok(None);
+    };
+
+    let contents = entry.contents(imfs)?;
+    let contents_str = str::from_utf8(contents)
+        .expect("File content was not valid UTF-8")
+        .to_string();
+
+    let properties = hashmap! {
+        "Source".to_owned() => RbxValue::String {
+            value: contents_str,
+        },
+    };
+
+    Ok(Some(InstanceSnapshot {
+        snapshot_id: None,
+        name: Cow::Owned(instance_name.to_owned()),
+        class_name: Cow::Borrowed(class_name),
+        properties,
+        children: Vec::new(),
+    }))
+}
+
+/// Attempts to snapshot an 'init' Lua script contained inside of a folder with
+/// the given name.
+///
+/// Scripts named `init.lua`, `init.server.lua`, or `init.client.lua` usurp
+/// their parents, which acts similarly to `__init__.py` from the Python world.
+fn snapshot_init<F: ImfsFetcher>(
+    imfs: &mut Imfs<F>,
+    folder_entry: &ImfsEntry,
+    init_name: &str,
+) -> SnapshotInstanceResult<'static> {
+    let init_path = folder_entry.path().join(init_name);
+
+    if let Some(init_entry) = imfs.get(init_path).with_not_found()? {
+        if let Some(mut dir_snapshot) = SnapshotDir::from_imfs(imfs, folder_entry)? {
+            if let Some(mut init_snapshot) = snapshot_lua_file(imfs, &init_entry)? {
+                init_snapshot.name = dir_snapshot.name;
+                init_snapshot.children = dir_snapshot.children;
+
+                return Ok(Some(init_snapshot));
+            }
+        }
+    }
+
+    Ok(None)
+}
+
 fn match_trailing<'a>(input: &'a str, trailer: &str) -> Option<&'a str> {
     if input.ends_with(trailer) {
         let end = input.len().saturating_sub(trailer.len());