Optimize table-like length storage

zap/src/config.rs

@@ -33,7 +33,7 @@ pub struct Config<'src> {
 	pub disable_fire_all: bool,
 }
 
-impl Config<'_> {
+impl<'src> Config<'src> {
 	pub fn server_reliable_count(&self) -> usize {
 		let reliable_count = self
 			.evdecls
@@ -75,6 +75,17 @@ impl Config<'_> {
 	pub fn client_reliable_ty(&self) -> NumTy {
 		NumTy::from_f64(0.0, self.client_reliable_count() as f64 - 1.0)
 	}
+
+	pub fn resolve_ty<'a>(&'a self, ty: &'a Ty<'src>) -> &'a Ty<'src> {
+		match ty {
+			Ty::Ref(name) => self
+				.tydecls
+				.iter()
+				.find(|decl| decl.name == *name)
+				.map_or(ty, |decl| &decl.ty),
+			_ => ty,
+		}
+	}
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -251,9 +262,9 @@ impl<'src> Ty<'src> {
 				}
 			}
 
-			Self::Map(..) => (2, None),
+			Self::Map(..) => (self.variants_size().unwrap().size(), None),
 
-			Self::Set(..) => (2, None),
+			Self::Set(..) => (self.variants_size().unwrap().size(), None),
 
 			Self::Opt(ty) => {
 				let (_, ty_max) = ty.size(tydecls, recursed);
@@ -313,6 +324,24 @@ impl<'src> Ty<'src> {
 			Self::Unknown => (0, None),
 		}
 	}
+
+	pub fn variants_size(&self) -> Option<NumTy> {
+		match self {
+			Ty::Enum(Enum::Unit(variants)) => Some(NumTy::from_f64(0.0, (variants.len() - 1) as f64)),
+			Ty::Enum(Enum::Tagged { variants, .. }) => Some(NumTy::from_f64(0.0, (variants.len() - 1) as f64)),
+			Ty::Num(numty, ..) => match numty {
+				NumTy::F32 => None,
+				NumTy::F64 => None,
+				NumTy::U8 => Some(NumTy::U8),
+				NumTy::U16 => Some(NumTy::U16),
+				NumTy::U32 => None,
+				NumTy::I8 => Some(NumTy::U8),
+				NumTy::I16 => Some(NumTy::U16),
+				NumTy::I32 => None,
+			},
+			_ => None,
+		}
+	}
 }
 
 #[derive(Debug, Clone)]

zap/src/irgen/des.rs

@@ -1,4 +1,4 @@
-use crate::config::{Enum, NumTy, Struct, Ty};
+use crate::config::{Config, Enum, NumTy, Struct, Ty};
 use std::collections::HashMap;
 
 use super::{Expr, Gen, Stmt, Var};
@@ -7,6 +7,7 @@ struct Des<'src> {
 	checks: bool,
 	buf: Vec<Stmt>,
 	var_occurrences: &'src mut HashMap<String, usize>,
+	config: &'src Config<'src>,
 }
 
 impl Gen for Des<'_> {
@@ -183,12 +184,14 @@ impl Des<'_> {
 			}
 
 			Ty::Map(key, val) => {
+				let length_numty = self.config.resolve_ty(key).variants_size().unwrap_or(NumTy::U16);
+
 				self.push_assign(into.clone(), Expr::EmptyTable);
 
 				self.push_stmt(Stmt::NumFor {
 					var: "_".into(),
 					from: 1.0.into(),
-					to: self.readu16(),
+					to: self.readnumty(length_numty),
 				});
 
 				let (key_name, key_expr) = self.add_occurrence("key");
@@ -205,12 +208,14 @@ impl Des<'_> {
 			}
 
 			Ty::Set(key) => {
+				let length_numty = self.config.resolve_ty(key).variants_size().unwrap_or(NumTy::U16);
+
 				self.push_assign(into.clone(), Expr::EmptyTable);
 
 				self.push_stmt(Stmt::NumFor {
 					var: "_".into(),
 					from: 1.0.into(),
-					to: self.readu16(),
+					to: self.readnumty(length_numty),
 				});
 
 				let (key_name, key_expr) = self.add_occurrence("key");
@@ -475,14 +480,21 @@ impl Des<'_> {
 	}
 }
 
-pub fn gen<'a, I>(types: I, names: &[String], checks: bool, var_occurrences: &mut HashMap<String, usize>) -> Vec<Stmt>
+pub fn gen<'a, I>(
+	types: I,
+	names: &[String],
+	checks: bool,
+	var_occurrences: &mut HashMap<String, usize>,
+	config: &Config<'_>,
+) -> Vec<Stmt>
 where
 	I: IntoIterator<Item = &'a Ty<'a>>,
 {
 	Des {
 		checks,
 		buf: vec![],
 		var_occurrences,
+		config,
 	}
 	.gen(names, types.into_iter())
 }

zap/src/irgen/ser.rs

@@ -1,4 +1,4 @@
-use crate::config::{Enum, NumTy, Struct, Ty};
+use crate::config::{Config, Enum, NumTy, Struct, Ty};
 use std::collections::HashMap;
 
 use super::{Expr, Gen, Stmt, Var};
@@ -7,6 +7,7 @@ struct Ser<'src> {
 	checks: bool,
 	buf: Vec<Stmt>,
 	var_occurrences: &'src mut HashMap<String, usize>,
+	config: &'src Config<'src>,
 }
 
 impl Gen for Ser<'_> {
@@ -187,7 +188,12 @@ impl Ser<'_> {
 				let (len_name, len_expr) = self.add_occurrence("len");
 				let (len_pos_name, len_pos_expr) = self.add_occurrence("len_pos");
 
-				self.push_local(len_pos_name.clone(), Some(Var::from("alloc").call(vec![2.0.into()])));
+				let length_size = self.config.resolve_ty(key).variants_size().unwrap_or(NumTy::U16).size();
+
+				self.push_local(
+					len_pos_name.clone(),
+					Some(Var::from("alloc").call(vec![(length_size as f64).into()])),
+				);
 				self.push_local(len_name.clone(), Some(0.0.into()));
 
 				let (key_name, _) = self.add_occurrence("k");
@@ -206,7 +212,7 @@ impl Ser<'_> {
 				self.push_stmt(Stmt::End);
 
 				self.push_stmt(Stmt::Call(
-					Var::from("buffer").nindex("writeu16"),
+					Var::from("buffer").nindex(format!("writeu{}", length_size * 8)),
 					None,
 					vec!["outgoing_buff".into(), len_pos_expr.clone(), len_expr.clone()],
 				));
@@ -216,7 +222,12 @@ impl Ser<'_> {
 				let (len_name, len_expr) = self.add_occurrence("len");
 				let (len_pos_name, len_pos_expr) = self.add_occurrence("len_pos");
 
-				self.push_local(len_pos_name.clone(), Some(Var::from("alloc").call(vec![2.0.into()])));
+				let length_size = self.config.resolve_ty(key).variants_size().unwrap_or(NumTy::U16).size();
+
+				self.push_local(
+					len_pos_name.clone(),
+					Some(Var::from("alloc").call(vec![(length_size as f64).into()])),
+				);
 				self.push_local(len_name.clone(), Some(0.0.into()));
 
 				let (key_name, _) = self.add_occurrence("k");
@@ -234,7 +245,7 @@ impl Ser<'_> {
 				self.push_stmt(Stmt::End);
 
 				self.push_stmt(Stmt::Call(
-					Var::from("buffer").nindex("writeu16"),
+					Var::from("buffer").nindex(format!("writeu{}", length_size * 8)),
 					None,
 					vec!["outgoing_buff".into(), len_pos_expr.clone(), len_expr.clone()],
 				));
@@ -400,14 +411,21 @@ impl Ser<'_> {
 	}
 }
 
-pub fn gen<'a, I>(types: I, names: &[String], checks: bool, var_occurrences: &mut HashMap<String, usize>) -> Vec<Stmt>
+pub fn gen<'a, I>(
+	types: I,
+	names: &[String],
+	checks: bool,
+	var_occurrences: &mut HashMap<String, usize>,
+	config: &Config<'_>,
+) -> Vec<Stmt>
 where
 	I: IntoIterator<Item = &'a Ty<'a>>,
 {
 	Ser {
 		checks,
 		buf: vec![],
 		var_occurrences,
+		config,
 	}
 	.gen(names, types.into_iter())
 }

zap/src/output/luau/client.rs

@@ -150,6 +150,7 @@ impl<'src> ClientOutput<'src> {
 			&["value".to_string()],
 			self.config.write_checks,
 			&mut HashMap::new(),
+			self.config,
 		);
 		self.push_stmts(statements);
 		self.dedent();
@@ -158,7 +159,13 @@ impl<'src> ClientOutput<'src> {
 		self.push_line(&format!("function types.read_{name}()"));
 		self.indent();
 		self.push_line("local value;");
-		let statements = &des::gen(&[ty.clone()], &["value".to_string()], false, &mut HashMap::new());
+		let statements = &des::gen(
+			&[ty.clone()],
+			&["value".to_string()],
+			false,
+			&mut HashMap::new(),
+			self.config,
+		);
 		self.push_stmts(statements);
 		self.push_line("return value");
 		self.dedent();
@@ -357,6 +364,7 @@ impl<'src> ClientOutput<'src> {
 				&get_unnamed_values("value", ev.data.len()),
 				true,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -461,6 +469,7 @@ impl<'src> ClientOutput<'src> {
 				&get_unnamed_values("value", data.len()),
 				true,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -559,6 +568,7 @@ impl<'src> ClientOutput<'src> {
 				&get_unnamed_values("value", ev.data.len()),
 				self.config.write_checks,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -766,6 +776,7 @@ impl<'src> ClientOutput<'src> {
 				&get_named_values(value, &ev.data),
 				self.config.write_checks,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -1153,6 +1164,7 @@ impl<'src> ClientOutput<'src> {
 					&get_named_values(value, &fndecl.args),
 					self.config.write_checks,
 					&mut self.var_occurrences,
+					self.config,
 				);
 				self.push_stmts(statements);
 			}

zap/src/output/luau/server.rs

@@ -133,6 +133,7 @@ impl<'a> ServerOutput<'a> {
 			&["value".to_string()],
 			self.config.write_checks,
 			&mut HashMap::new(),
+			self.config,
 		);
 		self.push_stmts(statements);
 		self.dedent();
@@ -141,7 +142,13 @@ impl<'a> ServerOutput<'a> {
 		self.push_line(&format!("function types.read_{name}()"));
 		self.indent();
 		self.push_line("local value;");
-		let statements = &des::gen(&[ty.clone()], &["value".to_string()], true, &mut HashMap::new());
+		let statements = &des::gen(
+			&[ty.clone()],
+			&["value".to_string()],
+			true,
+			&mut HashMap::new(),
+			self.config,
+		);
 		self.push_stmts(statements);
 		self.push_line("return value");
 		self.dedent();
@@ -345,6 +352,7 @@ impl<'a> ServerOutput<'a> {
 				&get_unnamed_values("value", ev.data.len()),
 				true,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -405,6 +413,7 @@ impl<'a> ServerOutput<'a> {
 				&get_unnamed_values("value", fndecl.args.len()),
 				true,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -450,7 +459,13 @@ impl<'a> ServerOutput<'a> {
 
 			if let Some(types) = &fndecl.rets {
 				let names: Vec<String> = (0..types.len()).map(|i| format!("ret_{}", i + 1)).collect();
-				let statements = &ser::gen(types, &names, self.config.write_checks, &mut self.var_occurrences);
+				let statements = &ser::gen(
+					types,
+					&names,
+					self.config.write_checks,
+					&mut self.var_occurrences,
+					self.config,
+				);
 				self.push_stmts(statements);
 			}
 
@@ -471,7 +486,13 @@ impl<'a> ServerOutput<'a> {
 
 			if let Some(types) = &fndecl.rets {
 				let names: Vec<String> = (0..types.len()).map(|i| format!("ret_{}", i + 1)).collect();
-				let statements = &ser::gen(types, &names, self.config.write_checks, &mut self.var_occurrences);
+				let statements = &ser::gen(
+					types,
+					&names,
+					self.config.write_checks,
+					&mut self.var_occurrences,
+					self.config,
+				);
 				self.push_stmts(statements);
 			}
 
@@ -542,6 +563,7 @@ impl<'a> ServerOutput<'a> {
 				&get_unnamed_values("value", ev.data.len()),
 				true,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -668,6 +690,7 @@ impl<'a> ServerOutput<'a> {
 				&get_named_values(value, parameters),
 				self.config.write_checks,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -714,6 +737,7 @@ impl<'a> ServerOutput<'a> {
 				&get_named_values(value, parameters),
 				self.config.write_checks,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -774,6 +798,7 @@ impl<'a> ServerOutput<'a> {
 				&get_named_values(value, parameters),
 				self.config.write_checks,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -846,6 +871,7 @@ impl<'a> ServerOutput<'a> {
 				&get_named_values(value, parameters),
 				self.config.write_checks,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}
@@ -910,6 +936,7 @@ impl<'a> ServerOutput<'a> {
 				&get_named_values(value, parameters),
 				self.config.write_checks,
 				&mut self.var_occurrences,
+				self.config,
 			);
 			self.push_stmts(statements);
 		}