Support extension variables in record and tag union encoding derive keys

crates/compiler/derive_key/src/encoding.rs

@@ -2,10 +2,10 @@ use roc_module::{
     ident::{Lowercase, TagName},
     symbol::Symbol,
 };
-use roc_types::subs::{Content, FlatType, GetSubsSlice, Subs, Variable};
+use roc_types::subs::{Content, FlatType, Subs, Variable};
 
 use crate::{
-    util::{check_empty_ext_var, debug_name_record},
+    util::{check_derivable_ext_var, debug_name_record},
     DeriveError,
 };
 
@@ -63,12 +63,17 @@ impl FlatEncodable {
                     _ => Err(Underivable),
                 },
                 FlatType::Record(fields, ext) => {
-                    check_empty_ext_var(subs, ext, |ext| {
+                    let (fields_iter, ext) = fields.unsorted_iterator_and_ext(subs, ext);
+
+                    check_derivable_ext_var(subs, ext, |ext| {
                         matches!(ext, Content::Structure(FlatType::EmptyRecord))
                     })?;
 
-                    let mut field_names: Vec<_> =
+                    let mut field_names = Vec::with_capacity(fields.len());
-                        subs.get_subs_slice(fields.field_names()).to_vec();
+                    for (field_name, _) in fields_iter {
+                        field_names.push(field_name.clone());
+                    }
+
                     field_names.sort();
 
                     Ok(Key(FlatEncodableKey::Record(field_names)))
@@ -83,20 +88,23 @@ impl FlatEncodable {
                     //   [ A t1, B t1 t2 ] as R
                     // look the same on the surface, because `R` is only somewhere inside of the
                     // `t`-prefixed payload types.
-                    check_empty_ext_var(subs, ext, |ext| {
+                    let (tags_iter, ext) = tags.unsorted_tags_and_ext(subs, ext);
+
+                    check_derivable_ext_var(subs, ext, |ext| {
                         matches!(ext, Content::Structure(FlatType::EmptyTagUnion))
                     })?;
 
-                    let mut tag_names_and_payload_sizes: Vec<_> = tags
+                    let mut tag_names_and_payload_sizes: Vec<_> = tags_iter
-                        .iter_all()
+                        .tags
-                        .map(|(name_index, payload_slice_index)| {
+                        .into_iter()
-                            let payload_slice = subs[payload_slice_index];
+                        .map(|(name, payload_slice)| {
-                            let payload_size = payload_slice.length;
+                            let payload_size = payload_slice.len();
-                            let name = &subs[name_index];
+                            (name.clone(), payload_size as _)
-                            (name.clone(), payload_size)
                         })
                         .collect();
+
                     tag_names_and_payload_sizes.sort_by(|(t1, _), (t2, _)| t1.cmp(t2));
+
                     Ok(Key(FlatEncodableKey::TagUnion(tag_names_and_payload_sizes)))
                 }
                 FlatType::FunctionOrTagUnion(name_index, _, _) => Ok(Key(

crates/compiler/derive_key/src/util.rs

@@ -3,13 +3,15 @@ use roc_types::subs::{Content, Subs, Variable};
 
 use crate::DeriveError;
 
-pub(crate) fn check_empty_ext_var(
+pub(crate) fn check_derivable_ext_var(
     subs: &Subs,
     ext_var: Variable,
     is_empty_ext: impl Fn(&Content) -> bool,
 ) -> Result<(), DeriveError> {
     let ext_content = subs.get_content_without_compacting(ext_var);
-    if is_empty_ext(ext_content) {
+    if is_empty_ext(ext_content)
+        || matches!(ext_content, Content::FlexVar(_) | Content::RigidVar(_))
+    {
         Ok(())
     } else {
         match ext_content {

crates/compiler/test_derive/src/encoding.rs

@@ -8,10 +8,10 @@ use insta::assert_snapshot;
 
 use crate::{
     test_key_eq, test_key_neq,
-    util::{check_immediate, derive_test},
+    util::{check_derivable, check_immediate, derive_test},
     v,
 };
-use roc_derive_key::DeriveBuiltin::ToEncoder;
+use roc_derive_key::{encoding::FlatEncodableKey, DeriveBuiltin::ToEncoder, DeriveKey};
 use roc_module::symbol::Symbol;
 use roc_types::subs::Variable;
 
@@ -118,6 +118,45 @@ fn immediates() {
     check_immediate(ToEncoder, v!(STR), Symbol::ENCODE_STRING);
 }
 
+#[test]
+fn derivable_record_ext_flex_var() {
+    check_derivable(
+        ToEncoder,
+        v!({ a: v!(STR), }* ),
+        DeriveKey::ToEncoder(FlatEncodableKey::Record(vec!["a".into()])),
+    );
+}
+
+#[test]
+fn derivable_record_with_record_ext() {
+    check_derivable(
+        ToEncoder,
+        v!({ b: v!(STR), }{ a: v!(STR), } ),
+        DeriveKey::ToEncoder(FlatEncodableKey::Record(vec!["a".into(), "b".into()])),
+    );
+}
+
+#[test]
+fn derivable_tag_ext_flex_var() {
+    check_derivable(
+        ToEncoder,
+        v!([ A v!(STR) ]* ),
+        DeriveKey::ToEncoder(FlatEncodableKey::TagUnion(vec![("A".into(), 1)])),
+    );
+}
+
+#[test]
+fn derivable_tag_with_tag_ext() {
+    check_derivable(
+        ToEncoder,
+        v!([ B v!(STR) v!(U8) ][ A v!(STR) ]),
+        DeriveKey::ToEncoder(FlatEncodableKey::TagUnion(vec![
+            ("A".into(), 1),
+            ("B".into(), 2),
+        ])),
+    );
+}
+
 #[test]
 fn empty_record() {
     derive_test(ToEncoder, v!(EMPTY_RECORD), |golden| {

crates/compiler/test_derive/src/util.rs

@@ -56,7 +56,7 @@ fn module_source_and_path(builtin: DeriveBuiltin) -> (ModuleId, &'static str, Pa
 /// DSL for creating [`Content`][roc_types::subs::Content].
 #[macro_export]
 macro_rules! v {
-     ({ $($field:ident: $make_v:expr,)* $(?$opt_field:ident : $make_opt_v:expr,)* }) => {{
+     ({ $($field:ident: $make_v:expr,)* $(?$opt_field:ident : $make_opt_v:expr,)* }$( $ext:tt )?) => {{
          #[allow(unused)]
          use roc_types::types::RecordField;
          use roc_types::subs::{Subs, RecordFields, Content, FlatType, Variable};
@@ -68,7 +68,12 @@ macro_rules! v {
                  $( (stringify!($opt_field).into(), RecordField::Optional($opt_field)) ,)*
              ];
              let fields = RecordFields::insert_into_subs(subs, fields);
-             roc_derive::synth_var(subs, Content::Structure(FlatType::Record(fields, Variable::EMPTY_RECORD)))
+
+             #[allow(unused_mut, unused)]
+             let mut ext = Variable::EMPTY_RECORD;
+             $( ext = $crate::v!($ext)(subs); )?
+
+             roc_derive::synth_var(subs, Content::Structure(FlatType::Record(fields, ext)))
          }
      }};
      ([ $($tag:ident $($payload:expr)*),* ]$( $ext:tt )?) => {{
@@ -82,7 +87,7 @@ macro_rules! v {
              )*
              let tags = UnionTags::insert_into_subs::<_, Vec<Variable>>(subs, vec![ $( (TagName(stringify!($tag).into()), $tag) ,)* ]);
 
-             #[allow(unused_mut)]
+             #[allow(unused_mut, unused)]
              let mut ext = Variable::EMPTY_TAG_UNION;
              $( ext = $crate::v!($ext)(subs); )?
 
@@ -197,6 +202,18 @@ macro_rules! test_key_neq {
     )*};
 }
 
+pub(crate) fn check_derivable<Sy>(builtin: DeriveBuiltin, synth: Sy, key: DeriveKey)
+where
+    Sy: FnOnce(&mut Subs) -> Variable,
+{
+    let mut subs = Subs::new();
+    let var = synth(&mut subs);
+
+    let derived = Derived::builtin(builtin, &subs, var);
+
+    assert_eq!(derived, Ok(Derived::Key(key)));
+}
+
 pub(crate) fn check_underivable<Sy>(builtin: DeriveBuiltin, synth: Sy, err: DeriveError)
 where
     Sy: FnOnce(&mut Subs) -> Variable,