From f3185b3874e453bf343e9d42eb2f5507ee24fa20 Mon Sep 17 00:00:00 2001
From: Arthur Cohen <arthur.cohen@embecosm.com>
Date: Thu, 29 Aug 2024 15:30:42 +0200
Subject: [PATCH] widen: Keep track of nodes when creating new TypeMap

---
 typecheck/src/typemap.rs | 4 ++--
 typecheck/src/widen.rs   | 6 +++++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/typecheck/src/typemap.rs b/typecheck/src/typemap.rs
index 98381b92..715444b6 100644
--- a/typecheck/src/typemap.rs
+++ b/typecheck/src/typemap.rs
@@ -23,8 +23,8 @@ pub struct TypeRef(pub(crate) OriginIdx); // FIXME: Remove vis?
 
 #[derive(Debug)]
 pub struct TypeMap {
-    nodes: HashMap<OriginIdx, TypeRef>,
-    // FIXME: Remove?
+    // FIXME: Remove pub-ability?
+    pub(crate) nodes: HashMap<OriginIdx, TypeRef>,
     pub(crate) types: HashMap<TypeRef, Type>,
 }
 
diff --git a/typecheck/src/widen.rs b/typecheck/src/widen.rs
index 6ed689a9..873aef7b 100644
--- a/typecheck/src/widen.rs
+++ b/typecheck/src/widen.rs
@@ -33,7 +33,7 @@ fn widen_inner(
 }
 
 pub fn widen(fir: &Fir<FlattenData<'_>>, type_ctx: TypeCtx<TypeMap>) -> TypeCtx<TypeMap> {
-    let types = type_ctx
+    let mut types = type_ctx
         .types
         .types
         .iter()
@@ -48,6 +48,10 @@ pub fn widen(fir: &Fir<FlattenData<'_>>, type_ctx: TypeCtx<TypeMap>) -> TypeCtx<
             tymap
         });
 
+    // Don't forget to keep track of all the nodes that make use of the types
+    // we just widened
+    types.nodes = type_ctx.types.nodes;
+
     TypeCtx {
         primitives: type_ctx.primitives,
         types,