From 5622aa661e03409b6565b069ab8dec79cb8dc962 Mon Sep 17 00:00:00 2001
From: Rerumu <25379555+Rerumu@users.noreply.github.com>
Date: Sat, 24 Jun 2023 19:05:00 -0400
Subject: [PATCH] Add memory index support in analysis

---
 codegen/luajit/src/analyzer/localize.rs | 18 +++++++++++++++---
 codegen/luau/src/analyzer/localize.rs   | 18 +++++++++++++++---
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/codegen/luajit/src/analyzer/localize.rs b/codegen/luajit/src/analyzer/localize.rs
index 3c33e4b..4701f65 100644
--- a/codegen/luajit/src/analyzer/localize.rs
+++ b/codegen/luajit/src/analyzer/localize.rs
@@ -1,7 +1,10 @@
 use std::collections::BTreeSet;
 
 use wasm_ast::{
-	node::{BinOp, CmpOp, FuncData, LoadAt, MemoryGrow, MemorySize, StoreAt, UnOp},
+	node::{
+		BinOp, CmpOp, FuncData, LoadAt, MemoryCopy, MemoryFill, MemoryGrow, MemorySize, StoreAt,
+		UnOp,
+	},
 	visit::{Driver, Visitor},
 };
 
@@ -16,14 +19,14 @@ impl Visitor for Visit {
 	fn visit_load_at(&mut self, v: &LoadAt) {
 		let name = v.load_type().as_name();
 
-		self.memory_set.insert(0);
+		self.memory_set.insert(v.memory());
 		self.local_set.insert(("load", name));
 	}
 
 	fn visit_store_at(&mut self, v: &StoreAt) {
 		let name = v.store_type().as_name();
 
-		self.memory_set.insert(0);
+		self.memory_set.insert(v.memory());
 		self.local_set.insert(("store", name));
 	}
 
@@ -60,6 +63,15 @@ impl Visitor for Visit {
 	fn visit_memory_grow(&mut self, m: &MemoryGrow) {
 		self.memory_set.insert(m.memory());
 	}
+
+	fn visit_memory_copy(&mut self, m: &MemoryCopy) {
+		self.memory_set.insert(m.destination().memory());
+		self.memory_set.insert(m.source().memory());
+	}
+
+	fn visit_memory_fill(&mut self, m: &MemoryFill) {
+		self.memory_set.insert(m.destination().memory());
+	}
 }
 
 pub fn visit(ast: &FuncData) -> (BTreeSet<(&'static str, &'static str)>, BTreeSet<usize>) {
diff --git a/codegen/luau/src/analyzer/localize.rs b/codegen/luau/src/analyzer/localize.rs
index 9c1c6dd..6755072 100644
--- a/codegen/luau/src/analyzer/localize.rs
+++ b/codegen/luau/src/analyzer/localize.rs
@@ -1,7 +1,10 @@
 use std::collections::BTreeSet;
 
 use wasm_ast::{
-	node::{BinOp, CmpOp, FuncData, LoadAt, MemoryGrow, MemorySize, StoreAt, UnOp, Value},
+	node::{
+		BinOp, CmpOp, FuncData, LoadAt, MemoryCopy, MemoryFill, MemoryGrow, MemorySize, StoreAt,
+		UnOp, Value,
+	},
 	visit::{Driver, Visitor},
 };
 use wasmparser::ValType;
@@ -17,14 +20,14 @@ impl Visitor for Visit {
 	fn visit_load_at(&mut self, v: &LoadAt) {
 		let name = v.load_type().as_name();
 
-		self.memory_set.insert(0);
+		self.memory_set.insert(v.memory());
 		self.local_set.insert(("load", name));
 	}
 
 	fn visit_store_at(&mut self, v: &StoreAt) {
 		let name = v.store_type().as_name();
 
-		self.memory_set.insert(0);
+		self.memory_set.insert(v.memory());
 		self.local_set.insert(("store", name));
 	}
 
@@ -72,6 +75,15 @@ impl Visitor for Visit {
 	fn visit_memory_grow(&mut self, m: &MemoryGrow) {
 		self.memory_set.insert(m.memory());
 	}
+
+	fn visit_memory_copy(&mut self, m: &MemoryCopy) {
+		self.memory_set.insert(m.destination().memory());
+		self.memory_set.insert(m.source().memory());
+	}
+
+	fn visit_memory_fill(&mut self, m: &MemoryFill) {
+		self.memory_set.insert(m.destination().memory());
+	}
 }
 
 pub fn visit(ast: &FuncData) -> (BTreeSet<(&'static str, &'static str)>, BTreeSet<usize>) {