From 7c0385f86607fe4146da719e12d81b1bcaf0f46a Mon Sep 17 00:00:00 2001
From: Arya Irani <arya@lrrr.home>
Date: Thu, 5 Dec 2024 19:02:17 -0500
Subject: [PATCH] `run` should use the `TypeLookup` from the file too

fixes #5448
---
 parser-typechecker/src/Unison/UnisonFile.hs        | 10 ++++++++++
 .../src/Unison/Codebase/Editor/HandleInput/Run.hs  | 14 ++++++++------
 unison-src/transcripts/idempotent/fix5448.md       | 12 ++++++++++++
 3 files changed, 30 insertions(+), 6 deletions(-)
 create mode 100644 unison-src/transcripts/idempotent/fix5448.md

diff --git a/parser-typechecker/src/Unison/UnisonFile.hs b/parser-typechecker/src/Unison/UnisonFile.hs
index 785482bac6..b3b8a12e1d 100644
--- a/parser-typechecker/src/Unison/UnisonFile.hs
+++ b/parser-typechecker/src/Unison/UnisonFile.hs
@@ -33,6 +33,7 @@ module Unison.UnisonFile
     nonEmpty,
     termSignatureExternalLabeledDependencies,
     topLevelComponents,
+    typecheckedToTypeLookup,
     typecheckedUnisonFile,
     Unison.UnisonFile.rewrite,
     prepareRewrite,
@@ -368,6 +369,15 @@ declsToTypeLookup uf =
   where
     wrangle = Map.fromList . Map.elems
 
+typecheckedToTypeLookup :: TypecheckedUnisonFile v a -> TL.TypeLookup v a
+typecheckedToTypeLookup tuf =
+  TL.TypeLookup
+    mempty
+    (wrangle (dataDeclarations' tuf))
+    (wrangle (effectDeclarations' tuf))
+  where
+    wrangle = Map.fromList . Map.elems
+
 -- Returns true if the file has any definitions or watches
 nonEmpty :: TypecheckedUnisonFile v a -> Bool
 nonEmpty uf =
diff --git a/unison-cli/src/Unison/Codebase/Editor/HandleInput/Run.hs b/unison-cli/src/Unison/Codebase/Editor/HandleInput/Run.hs
index 9cf1cbeaff..53b1f42b9e 100644
--- a/unison-cli/src/Unison/Codebase/Editor/HandleInput/Run.hs
+++ b/unison-cli/src/Unison/Codebase/Editor/HandleInput/Run.hs
@@ -42,6 +42,7 @@ import Unison.UnisonFile (TypecheckedUnisonFile)
 import Unison.UnisonFile qualified as UF
 import Unison.UnisonFile.Names qualified as UF
 import Unison.Util.Defns (Defns (..))
+import Unison.Util.Monoid qualified as Monoid
 import Unison.Util.Recursion
 import Unison.Var qualified as Var
 
@@ -110,7 +111,7 @@ getTerm' mainName =
           mainToFile (MainTerm.BadType _ ty) = pure $ maybe NoTermWithThatName TermHasBadType ty
           mainToFile (MainTerm.Success hq tm typ) =
             let v = Var.named (HQ.toText hq)
-             in checkType typ \otyp ->
+             in checkType Nothing typ \otyp ->
                   pure (GetTermSuccess (v, tm, typ, otyp))
       getFromFile uf = do
         let components = join $ UF.topLevelComponents uf
@@ -118,21 +119,22 @@ getTerm' mainName =
         let mainComponent = filter ((\v -> Var.name v == HQ.toText mainName) . view _1) components
         case mainComponent of
           [(v, _, tm, ty)] ->
-            checkType ty \otyp ->
+            checkType (Just uf) ty \otyp ->
               let runMain = DD.forceTerm a a (Term.var a v)
                   v2 = Var.freshIn (Set.fromList [v]) v
                   a = ABT.annotation tm
                in pure (GetTermSuccess (v2, runMain, ty, otyp))
           _ -> getFromCodebase
-      checkType :: Type Symbol Ann -> (Type Symbol Ann -> Cli GetTermResult) -> Cli GetTermResult
-      checkType ty f = do
+      checkType :: Maybe (TypecheckedUnisonFile Symbol Ann) -> Type Symbol Ann -> (Type Symbol Ann -> Cli GetTermResult) -> Cli GetTermResult
+      checkType mayTuf ty f = do
         Cli.Env {codebase, runtime} <- ask
         case Typechecker.fitsScheme ty (Runtime.mainType runtime) of
           True -> do
-            typeLookup <-
+            tlCodebase <-
               Cli.runTransaction $
                 Codebase.typeLookupForDependencies codebase Defns {terms = Set.empty, types = Type.dependencies ty}
-            f $! synthesizeForce typeLookup ty
+            let tlTuf = Monoid.fromMaybe (fmap UF.typecheckedToTypeLookup mayTuf)
+            f $! synthesizeForce (tlTuf <> tlCodebase) ty
           False -> pure (TermHasBadType ty)
    in Cli.getLatestTypecheckedFile >>= \case
         Nothing -> getFromCodebase
diff --git a/unison-src/transcripts/idempotent/fix5448.md b/unison-src/transcripts/idempotent/fix5448.md
new file mode 100644
index 0000000000..fc71f75da7
--- /dev/null
+++ b/unison-src/transcripts/idempotent/fix5448.md
@@ -0,0 +1,12 @@
+``` unison :hide
+type NewType = NewType
+main = do NewType
+```
+
+You shouldn't have to `add` a type before using it with `run`.
+
+``` ucm
+scratch/main> run main
+
+  NewType
+```