@used_as_key dsl (#214)

* @used_as_key dsl

Allow marking a type as a key to auto-derive traits e.g. for utils code

Fixes #190

* @used_as_key test cases + update docs + recursive tagging + work for enums

docs/docs/comment_dsl.mdx

@@ -82,6 +82,20 @@ pub struct Bar {
 }
 ```
 
+## @used_as_key
+
+```cddl
+foo = [
+  x: uint,
+  y: uint,
+] ; @used_as_Key
+```
+
+cddl-codegen automatically derives `Ord`/`PartialOrd` or `Hash` for any types used within as a key in another type.
+Putting this comment on a type forces that type to derive those traits even if it weren't used in a key in the cddl spec.
+This is useful for when you are writing utility code that would put them in a map and want the generated code to have it already,
+which is particularly useful for re-generating as it lets your `mod.rs` files remain untouched.
+
 ## _CDDL_CODEGEN_EXTERN_TYPE_
 
 While not as a comment, this allows you to compose in hand-written structs into a cddl spec.

src/comment_ast.rs

@@ -11,6 +11,7 @@ pub struct RuleMetadata {
     pub name: Option<String>,
     pub is_newtype: bool,
     pub no_alias: bool,
+    pub used_as_key: bool,
 }
 
 pub fn merge_metadata(r1: &RuleMetadata, r2: &RuleMetadata) -> RuleMetadata {
@@ -24,6 +25,7 @@ pub fn merge_metadata(r1: &RuleMetadata, r2: &RuleMetadata) -> RuleMetadata {
         },
         is_newtype: r1.is_newtype || r2.is_newtype,
         no_alias: r1.no_alias || r2.no_alias,
+        used_as_key: r1.used_as_key || r2.used_as_key,
     };
     merged.verify();
     merged
@@ -33,6 +35,7 @@ enum ParseResult {
     NewType,
     Name(String),
     DontGenAlias,
+    UsedAsKey,
 }
 
 impl RuleMetadata {
@@ -54,6 +57,10 @@ impl RuleMetadata {
                 ParseResult::DontGenAlias => {
                     base.no_alias = true;
                 }
+
+                ParseResult::UsedAsKey => {
+                    base.used_as_key = true;
+                }
             }
         }
         base.verify();
@@ -88,9 +95,15 @@ fn tag_no_alias(input: &str) -> IResult<&str, ParseResult> {
     Ok((input, ParseResult::DontGenAlias))
 }
 
+fn tag_used_as_key(input: &str) -> IResult<&str, ParseResult> {
+    let (input, _) = tag("@used_as_key")(input)?;
+
+    Ok((input, ParseResult::UsedAsKey))
+}
+
 fn whitespace_then_tag(input: &str) -> IResult<&str, ParseResult> {
     let (input, _) = take_while(char::is_whitespace)(input)?;
-    let (input, result) = alt((tag_name, tag_newtype, tag_no_alias))(input)?;
+    let (input, result) = alt((tag_name, tag_newtype, tag_no_alias, tag_used_as_key))(input)?;
 
     Ok((input, result))
 }
@@ -130,6 +143,7 @@ fn parse_comment_name() {
                 name: Some("foo".to_string()),
                 is_newtype: false,
                 no_alias: false,
+                used_as_key: false,
             }
         ))
     );
@@ -145,6 +159,7 @@ fn parse_comment_newtype() {
                 name: None,
                 is_newtype: true,
                 no_alias: false,
+                used_as_key: false,
             }
         ))
     );
@@ -160,6 +175,39 @@ fn parse_comment_newtype_and_name() {
                 name: Some("foo".to_string()),
                 is_newtype: true,
                 no_alias: false,
+                used_as_key: false,
+            }
+        ))
+    );
+}
+
+#[test]
+fn parse_comment_newtype_and_name_and_used_as_key() {
+    assert_eq!(
+        rule_metadata("@newtype @used_as_key @name foo"),
+        Ok((
+            "",
+            RuleMetadata {
+                name: Some("foo".to_string()),
+                is_newtype: true,
+                no_alias: false,
+                used_as_key: true,
+            }
+        ))
+    );
+}
+
+#[test]
+fn parse_comment_used_as_key() {
+    assert_eq!(
+        rule_metadata("@used_as_key"),
+        Ok((
+            "",
+            RuleMetadata {
+                name: None,
+                is_newtype: false,
+                no_alias: false,
+                used_as_key: true,
             }
         ))
     );
@@ -175,6 +223,7 @@ fn parse_comment_newtype_and_name_inverse() {
                 name: Some("foo".to_string()),
                 is_newtype: true,
                 no_alias: false,
+                used_as_key: false,
             }
         ))
     );
@@ -190,6 +239,7 @@ fn parse_comment_name_noalias() {
                 name: Some("foo".to_string()),
                 is_newtype: false,
                 no_alias: true,
+                used_as_key: false,
             }
         ))
     );

src/intermediate.rs

@@ -539,6 +539,14 @@ impl<'a> IntermediateTypes<'a> {
                 k.visit_types(types, &mut |ty| mark_used_as_key(ty, used_as_key));
             }
         }
+        // do a recursive check on the ones explicitly tagged as keys using @used_as_key
+        // this is done here since the lambdas are defined here so we can reuse them
+        for ident in &self.used_as_key {
+            if let Some(rust_struct) = self.rust_struct(ident) {
+                rust_struct.visit_types(self, &mut |ty| mark_used_as_key(ty, &mut used_as_key));
+            }
+        }
+        // check all other places used as keys
         for rust_struct in self.rust_structs().values() {
             rust_struct.visit_types(self, &mut |ty| {
                 check_used_as_key(ty, self, &mut used_as_key)
@@ -547,7 +555,10 @@ impl<'a> IntermediateTypes<'a> {
                 domain.visit_types(self, &mut |ty| mark_used_as_key(ty, &mut used_as_key));
             }
         }
-        self.used_as_key = used_as_key;
+        // we use a separate one here to get around the borrow checker in the above visit_types
+        for ident in used_as_key {
+            self.mark_used_as_key(ident);
+        }
     }
 
     pub fn visit_types<F: FnMut(&ConceptualRustType)>(&self, f: &mut F) {
@@ -657,6 +668,10 @@ impl<'a> IntermediateTypes<'a> {
         self.used_as_key.contains(name)
     }
 
+    pub fn mark_used_as_key(&mut self, name: RustIdent) {
+        self.used_as_key.insert(name);
+    }
+
     pub fn print_info(&self) {
         if !self.plain_groups.is_empty() {
             println!("\n\nPlain groups:");

src/parsing.rs

@@ -173,6 +173,21 @@ fn parse_type_choices(
             AliasInfo::new(final_type, !rule_metadata.no_alias, !rule_metadata.no_alias),
         );
     } else {
+        let rule_metadata = merge_metadata(
+            &RuleMetadata::from(
+                type_choices
+                    .last()
+                    .and_then(|tc| tc.comments_after_type.as_ref()),
+            ),
+            &RuleMetadata::from(
+                type_choices
+                    .last()
+                    .and_then(|tc| tc.type1.comments_after_type.as_ref()),
+            ),
+        );
+        if rule_metadata.used_as_key {
+            types.mark_used_as_key(name.clone());
+        }
         let variants = create_variants_from_type_choices(types, parent_visitor, type_choices, cli);
         let rust_struct = RustStruct::new_type_choice(name.clone(), tag, variants, cli);
         match generic_params {
@@ -456,6 +471,9 @@ fn parse_type(
         &RuleMetadata::from(type1.comments_after_type.as_ref()),
         &RuleMetadata::from(type_choice.comments_after_type.as_ref()),
     );
+    if rule_metadata.used_as_key {
+        types.mark_used_as_key(type_name.clone());
+    }
     match &type1.type2 {
         Type2::Typename {
             ident,

tests/core/input.cddl

@@ -18,7 +18,7 @@ bar = {
 }
 
 plain = (d: #6.23(uint), e: tagged_text)
-outer = [a: uint, b: plain, c: "some text"]
+outer = [a: uint, b: plain, c: "some text"] ; @used_as_key
 plain_arrays = [
 ; this is not supported right now. When single-element arrays are supported remove this.
 ;	single: [plain],
@@ -34,7 +34,7 @@ table_arr_members = {
 
 c_enum = 3 / 1 / 4
 
-type_choice = 0 / "hello world" / uint / text / bytes / #6.64([*uint])
+type_choice = 0 / "hello world" / uint / text / bytes / #6.64([*uint]) ; @used_as_key
 
 non_overlapping_type_choice_all = uint / nint / text / bytes / #6.30("hello world") / [* uint] / { *text => uint }
 
@@ -45,7 +45,7 @@ enums = [
 	type_choice,
 ]
 
-group_choice = [ foo // 0, x: uint // plain ]
+group_choice = [ foo // 0, x: uint // plain ] ; @used_as_key
 
 foo_bytes = bytes .cbor foo
 

tests/core/tests.rs

@@ -370,4 +370,15 @@ mod tests {
         // b oob
         assert!(make_bounds(OOB::Lower, OOB::Upper, OOB::Lower, OOB::Upper, OOB::Upper, OOB::Above).is_err());
     }
+
+    #[test]
+    fn used_as_key() {
+        // this is just here to make sure this compiles (i.e. Ord traits are derived)
+        let mut set_outer: std::collections::BTreeSet<Outer> = std::collections::BTreeSet::new();
+        set_outer.insert(Outer::new(2143254, Plain::new(7576, String::from("wiorurri34h").into())));
+        let mut set_type_choice: std::collections::BTreeSet<TypeChoice> = std::collections::BTreeSet::new();
+        set_type_choice.insert(TypeChoice::Helloworld);
+        let mut set_group_choice: std::collections::BTreeSet<GroupChoice> = std::collections::BTreeSet::new();
+        set_group_choice.insert(GroupChoice::GroupChoice1(37));
+    }
 }

tests/preserve-encodings/input.cddl

@@ -1,4 +1,4 @@
-foo = #6.11([uint, text, bytes])
+foo = #6.11([uint, text, bytes]) ; @used_as_key
 
 bar = {
 	foo: #6.13(foo),
@@ -32,7 +32,7 @@ type_choice = 0 / "hello world" / uint / text / #6.16([*uint])
 
 non_overlapping_type_choice_all = uint / nint / text / bytes / #6.13("hello world") / [* uint] / { *text => uint }
 
-non_overlapping_type_choice_some = uint / nint / text
+non_overlapping_type_choice_some = uint / nint / text ; @used_as_key
 
 c_enum = 3 / 1 / 4
 

tests/preserve-encodings/tests.rs

@@ -792,4 +792,13 @@ mod tests {
         // b oob
         assert!(make_bounds(OOB::Lower, OOB::Upper, OOB::Lower, OOB::Upper, OOB::Upper, OOB::Above).is_err());
     }
+
+    #[test]
+    fn used_as_key() {
+        // this is just here to make sure this compiles (i.e. Hash/Eq traits are derived)
+        let mut set_foo: std::collections::HashSet<Foo> = std::collections::HashSet::new();
+        set_foo.insert(Foo::new(0, "text".to_owned(), vec![]));
+        let mut set_non_overlap: std::collections::HashSet<NonOverlappingTypeChoiceSome> = std::collections::HashSet::new();
+        set_non_overlap.insert(NonOverlappingTypeChoiceSome::new_uint(0));
+    }
 }