From 289301fc1791f10b67f702f91493bfa610aac393 Mon Sep 17 00:00:00 2001
From: Nick Gasson <nick@nickg.me.uk>
Date: Tue, 1 Aug 2023 16:22:47 +0100
Subject: [PATCH] Fix tests again

---
 src/lower.c    | 40 ++++++++++++++--------------------------
 src/lower.h    |  1 +
 src/rt/shell.c |  4 ++++
 3 files changed, 19 insertions(+), 26 deletions(-)

diff --git a/src/lower.c b/src/lower.c
index 90479065..721febb5 100644
--- a/src/lower.c
+++ b/src/lower.c
@@ -11482,7 +11482,7 @@ void unit_registry_purge(unit_registry_t *ur, ident_t prefix)
    }
 }
 
-static void unit_registry_put_all(unit_registry_t *ur, vcode_unit_t vu)
+void unit_registry_put_all(unit_registry_t *ur, vcode_unit_t vu)
 {
    vcode_select_unit(vu);
    unit_registry_put(ur, vcode_unit_name(), vu);
@@ -11505,37 +11505,25 @@ vcode_unit_t unit_registry_get(unit_registry_t *ur, ident_t ident)
 
       ident_t unit_name = ident_prefix(lname, uname, '.');
       tree_t unit = lib_get(lib, unit_name);
-      if (unit == NULL)
+      if (unit == NULL || !is_package(unit))
          return NULL;
-      else if (tree_kind(unit) == T_ENTITY) {
-         ident_t ename = ident_prefix(unit_name, well_known(W_ELAB), '.');
-         if ((unit = lib_get(lib, ename)) == NULL)
-            return NULL;
 
-         vcode_unit_t root = lib_get_vcode(lib, unit);
-         if (root == NULL)
+      if (tree_kind(unit) == T_PACKAGE && package_needs_body(unit)) {
+         tree_t body = body_of(unit);
+         if (body == NULL)
             return NULL;
 
-         unit_registry_put_all(ur, root);
+         unit_registry_defer(ur, unit_name, NULL, emit_package,
+                             lower_pack_body, tree_to_object(body));
       }
-      else if (is_package(unit)) {
-         if (tree_kind(unit) == T_PACKAGE && package_needs_body(unit)) {
-            tree_t body = body_of(unit);
-            if (body == NULL)
-               return NULL;
-
-            unit_registry_defer(ur, unit_name, NULL, emit_package,
-                                lower_pack_body, tree_to_object(body));
-         }
-         else
-            unit_registry_defer(ur, unit_name, NULL, emit_package,
-                                lower_package, tree_to_object(unit));
+      else
+         unit_registry_defer(ur, unit_name, NULL, emit_package,
+                             lower_package, tree_to_object(unit));
 
-         if (unit_name != ident) {
-            // We actually wanted a unit inside this package so need to
-            // force code generation
-            (void)unit_registry_get(ur, unit_name);
-         }
+      if (unit_name != ident) {
+         // We actually wanted a unit inside this package so need to
+         // force code generation
+         (void)unit_registry_get(ur, unit_name);
       }
 
       if ((ptr = hash_get(ur->map, ident)) == NULL)
diff --git a/src/lower.h b/src/lower.h
index cd8f5152..41349fee 100644
--- a/src/lower.h
+++ b/src/lower.h
@@ -34,6 +34,7 @@ void unit_registry_defer(unit_registry_t *ur, ident_t ident,
                          lower_fn_t fn, object_t *object);
 void unit_registry_purge(unit_registry_t *ur, ident_t prefix);
 bool unit_registry_query(unit_registry_t *ur, ident_t ident);
+void unit_registry_put_all(unit_registry_t *ur, vcode_unit_t vu);
 
 lower_unit_t *lower_unit_new(unit_registry_t *ur, lower_unit_t *parent,
                              vcode_unit_t vunit, cover_tagging_t *cover,
diff --git a/src/rt/shell.c b/src/rt/shell.c
index 7084ef27..6fae97e5 100644
--- a/src/rt/shell.c
+++ b/src/rt/shell.c
@@ -1173,6 +1173,10 @@ void shell_reset(tcl_shell_t *sh, tree_t top)
 
    sh->top = top;
 
+   vcode_unit_t vu = lib_get_vcode(lib_work(), top);
+   if (vu != NULL)
+     unit_registry_put_all(sh->registry, vu);
+
    shell_create_model(sh);
 
    sh->nsignals = count_signals(sh->root);
-- 
2.39.2