Consistent filtering in Rust client + minor fixes

crates/cli/src/subcommands/generate/rust.rs

@@ -1,21 +1,21 @@
 use super::code_indenter::CodeIndenter;
 use super::{GenCtx, GenItem};
 use convert_case::{Case, Casing};
-use spacetimedb_data_structures::map::HashSet;
 use spacetimedb_lib::sats::db::def::TableSchema;
 use spacetimedb_lib::sats::{
     AlgebraicType, AlgebraicTypeRef, ArrayType, BuiltinType, MapType, ProductType, ProductTypeElement, SumType,
     SumTypeVariant,
 };
 use spacetimedb_lib::{ReducerDef, TableDesc};
 use spacetimedb_primitives::ColList;
+use std::collections::BTreeSet;
 use std::fmt::{self, Write};
 use std::ops::Deref;
 
 type Indenter = CodeIndenter<String>;
 
 /// Pairs of (module_name, TypeName).
-type Imports = HashSet<(String, String)>;
+type Imports = BTreeSet<(String, String)>;
 
 enum MaybePrimitive<'a> {
     Primitive(&'a str),
@@ -470,7 +470,19 @@ fn print_table_filter_methods(ctx: &GenCtx, out: &mut Indenter, table_type_name:
         |out| {
             for field in table.columns() {
                 let field_name = field.col_name.deref().to_case(Case::Snake);
-                // TODO: ensure that fields are PartialEq
+                match &field.col_type {
+                    AlgebraicType::Product(prod) => {
+                        if !prod.is_special() {
+                            continue;
+                        }
+                    }
+                    AlgebraicType::Ref(_)
+                    | AlgebraicType::Sum(_)
+                    | AlgebraicType::Builtin(BuiltinType::Array(_) | BuiltinType::Map(_)) => {
+                        continue;
+                    }
+                    AlgebraicType::Builtin(_) => {}
+                }
                 writeln!(out, "{ALLOW_UNUSED}");
                 write!(out, "pub fn filter_by_{field_name}({field_name}: ");
                 // TODO: the filter methods should take the target value by
@@ -482,27 +494,31 @@ fn print_table_filter_methods(ctx: &GenCtx, out: &mut Indenter, table_type_name:
                 write!(out, ") -> ");
                 let ct = constraints[&ColList::new(field.col_pos)];
 
-                if ct.has_unique() {
-                    write!(out, "Option<Self>");
-                } else {
-                    write!(out, "TableIter<Self>");
-                }
+                write!(out, "TableIter<Self>");
                 out.delimited_block(
                     " {",
                     |out| {
                         writeln!(
                             out,
-                            "Self::{}(|row| row.{} == {})",
                             // TODO: for primary keys, we should be able to do better than
                             //       `find` or `filter`. We should be able to look up
                             //       directly in the `TableCache`.
-                            if ct.has_unique() { "find" } else { "filter" },
-                            field_name,
-                            field_name,
+                            "Self::filter(|row| row.{field_name} == {field_name})",
                         )
                     },
                     "}\n",
                 );
+                if ct.has_unique() {
+                    writeln!(out, "{ALLOW_UNUSED}");
+                    write!(out, "pub fn find_by_{field_name}({field_name}: ");
+                    write_type_ctx(ctx, out, &field.col_type);
+                    write!(out, ") -> Option<Self> ");
+                    out.delimited_block(
+                        "{",
+                        |out| writeln!(out, "Self::find(|row| row.{field_name} == {field_name})"),
+                        "}\n",
+                    );
+                }
             }
         },
         "}\n",
@@ -1106,7 +1122,7 @@ fn gen_and_print_imports<Roots, SearchFn>(
 ) where
     SearchFn: FnOnce(&GenCtx, &mut Imports, Roots),
 {
-    let mut imports = HashSet::new();
+    let mut imports = BTreeSet::new();
     search_fn(ctx, &mut imports, roots);
 
     print_imports(out, imports, this_file);

crates/cli/tests/codegen.rs

@@ -11,24 +11,25 @@ fn compiled_module() -> &'static Path {
         .path()
 }
 
-#[test]
-fn test_codegen_output() {
-    let outfiles: HashMap<_, _> = generate::generate(compiled_module(), generate::Language::Csharp, "SpacetimeDB")
-        .unwrap()
-        .into_iter()
-        .collect();
-    insta::with_settings!({ sort_maps => true }, {
-        insta::assert_toml_snapshot!(outfiles);
-    });
+macro_rules! declare_tests {
+    ($($name:ident => $lang:ident,)*) => {
+        $(
+            #[test]
+            fn $name() {
+                let outfiles: HashMap<_, _> = generate::generate(compiled_module(), generate::Language::$lang, "SpacetimeDB")
+                    .unwrap()
+                    .into_iter()
+                    .collect();
+                insta::with_settings!({ sort_maps => true }, {
+                    insta::assert_toml_snapshot!(outfiles);
+                });
+            }
+        )*
+    }
 }
 
-#[test]
-fn test_typescript_codegen_output() {
-    let outfiles: HashMap<_, _> = generate::generate(compiled_module(), generate::Language::TypeScript, "SpacetimeDB")
-        .unwrap()
-        .into_iter()
-        .collect();
-    insta::with_settings!({ sort_maps => true }, {
-        insta::assert_toml_snapshot!(outfiles);
-    });
+declare_tests! {
+    test_codegen_csharp => Csharp,
+    test_codegen_typescript => TypeScript,
+    test_codegen_rust => Rust,
 }