load union at index for recursive and nonrecursive unions

compiler/gen_dev/src/generic64/mod.rs

@@ -9,7 +9,7 @@ use roc_error_macros::internal_error;
 use roc_module::symbol::{Interns, Symbol};
 use roc_mono::code_gen_help::CodeGenHelp;
 use roc_mono::ir::{BranchInfo, JoinPointId, Literal, Param, ProcLayout, SelfRecursive, Stmt};
-use roc_mono::layout::{Builtin, Layout, UnionLayout};
+use roc_mono::layout::{Builtin, Layout, TagIdIntType, UnionLayout};
 use roc_target::TargetInfo;
 use std::marker::PhantomData;
 
@@ -879,6 +879,27 @@ impl<
             .load_field_at_index(sym, structure, index, field_layouts);
     }
 
+    fn load_union_at_index(
+        &mut self,
+        sym: &Symbol,
+        structure: &Symbol,
+        tag_id: TagIdIntType,
+        index: u64,
+        union_layout: &UnionLayout<'a>,
+    ) {
+        match union_layout {
+            UnionLayout::NonRecursive(tag_layouts) | UnionLayout::Recursive(tag_layouts) => {
+                self.storage_manager.load_field_at_index(
+                    sym,
+                    structure,
+                    index,
+                    tag_layouts[tag_id as usize],
+                );
+            }
+            x => todo!("loading from union type: {:?}", x),
+        }
+    }
+
     fn get_tag_id(&mut self, sym: &Symbol, structure: &Symbol, union_layout: &UnionLayout<'a>) {
         self.storage_manager
             .load_union_tag_id(&mut self.buf, sym, structure, union_layout);

compiler/gen_dev/src/lib.rs

@@ -14,7 +14,7 @@ use roc_mono::ir::{
     BranchInfo, CallType, Expr, JoinPointId, ListLiteralElement, Literal, Param, Proc, ProcLayout,
     SelfRecursive, Stmt,
 };
-use roc_mono::layout::{Builtin, Layout, LayoutId, LayoutIds, UnionLayout};
+use roc_mono::layout::{Builtin, Layout, LayoutId, LayoutIds, TagIdIntType, UnionLayout};
 
 mod generic64;
 mod object_builder;
@@ -315,6 +315,14 @@ trait Backend<'a> {
             } => {
                 self.load_struct_at_index(sym, structure, *index, field_layouts);
             }
+            Expr::UnionAtIndex {
+                structure,
+                tag_id,
+                union_layout,
+                index,
+            } => {
+                self.load_union_at_index(sym, structure, *tag_id, *index, union_layout);
+            }
             Expr::GetTagId {
                 structure,
                 union_layout,
@@ -686,6 +694,16 @@ trait Backend<'a> {
         field_layouts: &'a [Layout<'a>],
     );
 
+    /// load_union_at_index loads into `sym` the value at `index` for `tag_id`.
+    fn load_union_at_index(
+        &mut self,
+        sym: &Symbol,
+        structure: &Symbol,
+        tag_id: TagIdIntType,
+        index: u64,
+        union_layout: &UnionLayout<'a>,
+    );
+
     /// get_tag_id loads the tag id from a the union.
     fn get_tag_id(&mut self, sym: &Symbol, structure: &Symbol, union_layout: &UnionLayout<'a>);