Correct type of static forms in hsExprType

The simplest way to do this seemed to be to persist the whole type in
the extension field from the typechecker so that the few relevant places

 * Desugaring can work out the return type by splitting this type rather
   than calling `dsExpr` (slightly more efficient).
 * hsExprType can just return the correct type.
 * Zonking has to now zonk the type as well

The other option we considered was wiring in StaticPtr but that is
actually quite tricky because StaticPtr refers to StaticPtrInfo which
has field selectors (which we can't easily wire in).

Fixes #20150

12 files changed, 31 insertions, 10 deletions

diff --git a/compiler/GHC/Builtin/Types.hs b/compiler/GHC/Builtin/Types.hs
index 6be9ecd293..2096e27a2b 100644
--- a/compiler/GHC/Builtin/Types.hs
+++ b/compiler/GHC/Builtin/Types.hs
@@ -158,7 +158,6 @@ module GHC.Builtin.Types (
         naturalTy, naturalTyCon, naturalTyConName,
         naturalNSDataCon, naturalNSDataConName,
         naturalNBDataCon, naturalNBDataConName
-
     ) where
 
 import GHC.Prelude
diff --git a/compiler/GHC/Hs/Expr.hs b/compiler/GHC/Hs/Expr.hs
index 6228b7d90e..8c77966e18 100644
--- a/compiler/GHC/Hs/Expr.hs
+++ b/compiler/GHC/Hs/Expr.hs
@@ -351,7 +351,9 @@ type instance XProc          (GhcPass _) = EpAnn [AddEpAnn]
 
 type instance XStatic        GhcPs = EpAnn [AddEpAnn]
 type instance XStatic        GhcRn = NameSet
-type instance XStatic        GhcTc = NameSet
+type instance XStatic        GhcTc = (NameSet, Type)
+  -- Free variables and type of expression, this is stored for convenience as wiring in
+  -- StaticPtr is a bit tricky (see #20150)
 
 type instance XPragE         (GhcPass _) = NoExtField
 
diff --git a/compiler/GHC/Hs/Syn/Type.hs b/compiler/GHC/Hs/Syn/Type.hs
index 1501abbb9e..c985c9237c 100644
--- a/compiler/GHC/Hs/Syn/Type.hs
+++ b/compiler/GHC/Hs/Syn/Type.hs
@@ -138,7 +138,7 @@ hsExprType e@(HsSpliceE{}) = pprPanic "hsExprType: Unexpected HsSpliceE"
                                       -- can't use `dataConCantHappen` since they are still present before
                                       -- than in the typechecked AST.
 hsExprType (HsProc _ _ lcmd_top) = lhsCmdTopType lcmd_top
-hsExprType (HsStatic _ e) = lhsExprType e
+hsExprType (HsStatic (_, ty) _s) = ty
 hsExprType (HsPragE _ _ e) = lhsExprType e
 hsExprType (XExpr (WrapExpr (HsWrap wrap e))) = hsWrapperType wrap $ hsExprType e
 hsExprType (XExpr (ExpansionExpr (HsExpanded _ tc_e))) = hsExprType tc_e
diff --git a/compiler/GHC/HsToCore/Expr.hs b/compiler/GHC/HsToCore/Expr.hs
index c2501de165..4e4eca8cef 100644
--- a/compiler/GHC/HsToCore/Expr.hs
+++ b/compiler/GHC/HsToCore/Expr.hs
@@ -418,9 +418,9 @@ See Note [Grand plan for static forms] in GHC.Iface.Tidy.StaticPtrTable for an o
     g = ... makeStatic loc f ...
 -}
 
-dsExpr (HsStatic _ expr@(L loc _)) = do
+dsExpr (HsStatic (_, whole_ty) expr@(L loc _)) = do
     expr_ds <- dsLExpr expr
-    let ty = exprType expr_ds
+    let (_, [ty]) = splitTyConApp whole_ty
     makeStaticId <- dsLookupGlobalId makeStaticName
 
     dflags <- getDynFlags
diff --git a/compiler/GHC/Iface/Ext/Ast.hs b/compiler/GHC/Iface/Ext/Ast.hs
index 0c74db385d..83eb475a78 100644
--- a/compiler/GHC/Iface/Ext/Ast.hs
+++ b/compiler/GHC/Iface/Ext/Ast.hs
@@ -737,7 +737,6 @@ instance HiePass p => HasType (LocatedA (HsExpr (GhcPass p))) where
         HsLet _ _ _ _ body -> computeLType body
         RecordCon con_expr _ _ -> computeType con_expr
         ExprWithTySig _ e _ -> computeLType e
-        HsStatic _ e -> computeLType e
         HsPragE _ _ e -> computeLType e
         XExpr (ExpansionExpr (HsExpanded _ e)) -> computeType e
         XExpr (HsTick _ e) -> computeLType e
diff --git a/compiler/GHC/Tc/Gen/Expr.hs b/compiler/GHC/Tc/Gen/Expr.hs
index 0c1d4faf24..8bff4b7e53 100644
--- a/compiler/GHC/Tc/Gen/Expr.hs
+++ b/compiler/GHC/Tc/Gen/Expr.hs
@@ -455,9 +455,10 @@ tcExpr (HsStatic fvs expr) res_ty
                                              [p_ty]
         ; let wrap = mkWpEvVarApps [typeable_ev] <.> mkWpTyApps [expr_ty]
         ; loc <- getSrcSpanM
+        ; static_ptr_ty_con <- tcLookupTyCon staticPtrTyConName
         ; return $ mkHsWrapCo co $ HsApp noComments
                             (L (noAnnSrcSpan loc) $ mkHsWrap wrap fromStaticPtr)
-                            (L (noAnnSrcSpan loc) (HsStatic fvs expr'))
+                            (L (noAnnSrcSpan loc) (HsStatic (fvs, mkTyConApp static_ptr_ty_con [expr_ty]) expr'))
         }
 
 {-
diff --git a/compiler/GHC/Tc/Utils/Zonk.hs b/compiler/GHC/Tc/Utils/Zonk.hs
index fec8d90d5d..6a65d5d383 100644
--- a/compiler/GHC/Tc/Utils/Zonk.hs
+++ b/compiler/GHC/Tc/Utils/Zonk.hs
@@ -920,8 +920,9 @@ zonkExpr env (HsProc x pat body)
         ; return (HsProc x new_pat new_body) }
 
 -- StaticPointers extension
-zonkExpr env (HsStatic fvs expr)
-  = HsStatic fvs <$> zonkLExpr env expr
+zonkExpr env (HsStatic (fvs, ty) expr)
+  = do new_ty <- zonkTcTypeToTypeX env ty
+       HsStatic (fvs, new_ty) <$> zonkLExpr env expr
 
 zonkExpr env (XExpr (WrapExpr (HsWrap co_fn expr)))
   = do (env1, new_co_fn) <- zonkCoFn env co_fn
diff --git a/compiler/Language/Haskell/Syntax/Expr.hs b/compiler/Language/Haskell/Syntax/Expr.hs
index e6ce12f8ae..0baaeaa148 100644
--- a/compiler/Language/Haskell/Syntax/Expr.hs
+++ b/compiler/Language/Haskell/Syntax/Expr.hs
@@ -637,7 +637,7 @@ data HsExpr p
   -- | - 'GHC.Parser.Annotation.AnnKeywordId' : 'GHC.Parser.Annotation.AnnStatic',
 
   -- For details on above see note [exact print annotations] in GHC.Parser.Annotation
-  | HsStatic (XStatic p) -- Free variables of the body
+  | HsStatic (XStatic p) -- Free variables of the body, and type after typechecking
              (LHsExpr p)        -- Body
 
   ---------------------------------------
diff --git a/testsuite/tests/ghci/scripts/T20150.hs b/testsuite/tests/ghci/scripts/T20150.hs
new file mode 100644
index 0000000000..e1706dedc4
--- /dev/null
+++ b/testsuite/tests/ghci/scripts/T20150.hs
@@ -0,0 +1,9 @@
+{-# LANGUAGE StaticPointers #-}
+module T20150 where
+
+import GHC.StaticPtr
+
+foo :: StaticPtr Int
+foo = static 0
+
+
diff --git a/testsuite/tests/ghci/scripts/T20150.script b/testsuite/tests/ghci/scripts/T20150.script
new file mode 100644
index 0000000000..0b5d132cdd
--- /dev/null
+++ b/testsuite/tests/ghci/scripts/T20150.script
@@ -0,0 +1,3 @@
+:set +c
+:l T20150.hs
+:all-types
diff --git a/testsuite/tests/ghci/scripts/T20150.stdout b/testsuite/tests/ghci/scripts/T20150.stdout
new file mode 100644
index 0000000000..e55ee89ea9
--- /dev/null
+++ b/testsuite/tests/ghci/scripts/T20150.stdout
@@ -0,0 +1,6 @@
+Collecting type info for 1 module(s) ... 
+T20150.hs:(7,1)-(7,3): GHC.StaticPtr.StaticPtr GHC.Types.Int
+T20150.hs:(7,14)-(7,14): GHC.Types.Int
+T20150.hs:(7,7)-(7,14): GHC.StaticPtr.StaticPtr GHC.Types.Int
+T20150.hs:(7,7)-(7,14): GHC.StaticPtr.StaticPtr GHC.Types.Int -> GHC.StaticPtr.StaticPtr GHC.Types.Int
+T20150.hs:(7,7)-(7,14): GHC.StaticPtr.StaticPtr GHC.Types.Int
diff --git a/testsuite/tests/ghci/scripts/all.T b/testsuite/tests/ghci/scripts/all.T
index 5e9aea056d..71e0ea80a5 100755
--- a/testsuite/tests/ghci/scripts/all.T
+++ b/testsuite/tests/ghci/scripts/all.T
@@ -354,3 +354,4 @@ test('T20473b', normal, ghci_script, ['T20473b.script'])
 test('T20587', [extra_files(['../shell.hs'])], ghci_script,
                 ['T20587.script'])
 test('T20909', normal, ghci_script, ['T20909.script'])
+test('T20150', normal, ghci_script, ['T20150.script'])