Retrieving the datacon of an arbitrary closure

This patch extends the RTS linker and the dynamic linker so that it is possible to find out the datacon of a closure in heap at runtime:
- The RTS linker now carries a hashtable 'Address->Symbol' for data constructors
- The Persistent Linker State in the dynamic linker is extended in a similar way.

Finally, these two sources of information are consulted by:

> Linker.recoverDataCon :: a -> TcM Name

compiler/ghci/ByteCodeItbls.lhs

@@ -6,7 +6,7 @@ ByteCodeItbls: Generate infotables for interpreter-made bytecodes
 \begin{code}
 {-# OPTIONS -optc-DNON_POSIX_SOURCE #-}
 
-module ByteCodeItbls ( ItblEnv, ItblPtr(..), itblCode, mkITbls ) where
+module ByteCodeItbls ( ItblEnv, ItblPtr, mkITbls ) where
 
 #include "HsVersions.h"
 

compiler/ghci/ByteCodeLink.lhs

@@ -10,6 +10,7 @@ module ByteCodeLink (
 	HValue, 
 	ClosureEnv, emptyClosureEnv, extendClosureEnv,
 	linkBCO, lookupStaticPtr
+       ,lookupIE
   ) where
 
 #include "HsVersions.h"

compiler/ghci/Linker.lhs

@@ -18,6 +18,7 @@ module Linker ( HValue, showLinkerState,
 		linkExpr, unload, extendLinkEnv, withExtendedLinkEnv,
                 extendLoadedPkgs,
 		linkPackages,initDynLinker
+               ,recoverDataCon
 	) where
 
 #include "HsVersions.h"
@@ -26,7 +27,14 @@ import ObjLink
 import ByteCodeLink
 import ByteCodeItbls
 import ByteCodeAsm
-
+import RtClosureInspect
+import Var
+import IfaceEnv
+import Config
+import OccName
+import TcRnMonad
+import Constants
+import Encoding
 import Packages
 import DriverPhases
 import Finder
@@ -50,9 +58,12 @@ import SrcLoc
 
 -- Standard libraries
 import Control.Monad
- 
+import Control.Arrow    ( second )
+
 import Data.IORef
 import Data.List
+import Foreign.Ptr
+import GHC.Exts
 
 import System.IO
 import System.Directory
@@ -108,6 +119,7 @@ data PersistentLinkerState
 	-- Held, as usual, in dependency order; though I am not sure if
 	-- that is really important
 	pkgs_loaded :: [PackageId]
+       ,dtacons_env :: DataConEnv
      }
 
 emptyPLS :: DynFlags -> PersistentLinkerState
@@ -116,7 +128,9 @@ emptyPLS dflags = PersistentLinkerState {
 			itbl_env    = emptyNameEnv,
 			pkgs_loaded = init_pkgs,
 			bcos_loaded = [],
-			objs_loaded = [] }
+			objs_loaded = []
+                      , dtacons_env = emptyAddressEnv
+                                        }
   -- Packages that don't need loading, because the compiler 
   -- shares them with the interpreted program.
   --
@@ -138,6 +152,56 @@ extendLinkEnv new_bindings
 	    new_pls = pls { closure_env = new_closure_env }
 	writeIORef v_PersistentLinkerState new_pls
 
+
+recoverDataCon :: a -> TcM Name
+recoverDataCon a = recoverDCInRTS a `recoverM` ioToTcRn (do 
+               mb_name <- recoverDCInDynEnv a
+               maybe (fail "Linker.recoverDatacon: Name not found in Dyn Env")
+                     return
+                     mb_name)
+
+-- | If a is a Constr closure, lookupDC returns either the Name of the DataCon or the
+--   symbol if it is a nullary constructor
+--   For instance, for a closure containing 'Just x' it would return the Name for Data.Maybe.Just
+--   For a closure containing 'Nothing' it would return the String "DataziMaybe_Nothing_static_info"
+recoverDCInDynEnv :: a -> IO (Maybe Name)
+recoverDCInDynEnv a = do 
+   pls <- readIORef v_PersistentLinkerState
+   let de = dtacons_env pls
+   ctype <- getClosureType a
+   if not (isConstr ctype) 
+         then putStrLn ("Not a Constr (" ++ show  ctype ++ ")") >> 
+              return Nothing
+         else do let infot = getInfoTablePtr a
+                     name  = lookupAddressEnv de (castPtr$ infot `plusPtr` (wORD_SIZE*2))
+                 return name
+
+
+recoverDCInRTS :: a -> TcM Name 
+recoverDCInRTS a = do
+  ctype <- ioToTcRn$ getClosureType a
+  if (not$ isConstr ctype)
+     then fail "not Constr"
+     else do
+       Just symbol <- ioToTcRn$ lookupDataCon (getInfoTablePtr a)
+       let (occ,mod) = (parse . lex) symbol
+       lookupOrig mod occ
+    where lex x = map zDecodeString . init . init . split '_' . removeLeadingUnderscore $ x
+          parse [pkg, modName, occ] = (mkOccName OccName.dataName occ,
+              mkModule (stringToPackageId pkg) (mkModuleName modName))
+          parse [modName, occ] = (mkOccName OccName.dataName occ,
+              mkModule mainPackageId (mkModuleName modName))
+          split delim = let 
+                 helper [] = Nothing
+                 helper x  = Just . second (drop 1) . break (==delim) $ x
+              in unfoldr helper
+
+removeLeadingUnderscore = if cLeadingUnderscore=="YES" 
+                                       then tail 
+                                       else id
+
+
+
 withExtendedLinkEnv :: [(Name,HValue)] -> IO a -> IO a
 withExtendedLinkEnv new_env action
     = bracket set_new_env
@@ -173,7 +237,9 @@ showLinkerState
        printDump (vcat [text "----- Linker state -----",
 			text "Pkgs:" <+> ppr (pkgs_loaded pls),
 			text "Objs:" <+> ppr (objs_loaded pls),
-			text "BCOs:" <+> ppr (bcos_loaded pls)])
+			text "BCOs:" <+> ppr (bcos_loaded pls),
+                        text "DataCons:" <+> ppr (dtacons_env pls)
+                       ])
 \end{code}
 
 
@@ -324,6 +390,8 @@ linkExpr :: HscEnv -> SrcSpan -> UnlinkedBCO -> IO HValue
 --
 -- Raises an IO exception if it can't find a compiled version of the
 -- dependents to link.
+--
+-- Note: This function side-effects the linker state (Pepe)
 
 linkExpr hsc_env span root_ul_bco
   = do {  
@@ -353,9 +421,11 @@ linkExpr hsc_env span root_ul_bco
      pls <- readIORef v_PersistentLinkerState
    ; let ie = itbl_env pls
 	 ce = closure_env pls
+         de = dtacons_env pls
 
 	-- Link the necessary packages and linkables
-   ; (_, (root_hval:_)) <- linkSomeBCOs False ie ce [root_ul_bco]
+   ; (_,de_out, (root_hval:_)) <- linkSomeBCOs False ie ce de [root_ul_bco]
+   ; writeIORef v_PersistentLinkerState (pls{dtacons_env=de_out})
    ; return root_hval
    }}
    where
@@ -615,10 +685,11 @@ dynLinkBCOs bcos
 	    gce	      = closure_env pls
             final_ie  = foldr plusNameEnv (itbl_env pls) ies
 
-        (final_gce, linked_bcos) <- linkSomeBCOs True final_ie gce ul_bcos
+        (final_gce, final_de, linked_bcos) <- linkSomeBCOs True final_ie gce (dtacons_env pls) ul_bcos
 		-- What happens to these linked_bcos?
 
 	let pls2 = pls1 { closure_env = final_gce,
+                          dtacons_env = final_de, 
 			  itbl_env    = final_ie }
 
 	writeIORef v_PersistentLinkerState pls2
@@ -629,19 +700,18 @@ linkSomeBCOs :: Bool 	-- False <=> add _all_ BCOs to returned closure env
                         -- True  <=> add only toplevel BCOs to closure env
              -> ItblEnv 
              -> ClosureEnv 
+             -> DataConEnv
              -> [UnlinkedBCO]
-             -> IO (ClosureEnv, [HValue])
+             -> IO (ClosureEnv, DataConEnv, [HValue])
 			-- The returned HValues are associated 1-1 with
 			-- the incoming unlinked BCOs.  Each gives the
 			-- value of the corresponding unlinked BCO
 					
-
-linkSomeBCOs toplevs_only ie ce_in ul_bcos
+linkSomeBCOs toplevs_only ie ce_in de_in ul_bcos
    = do let nms = map unlinkedBCOName ul_bcos
         hvals <- fixIO 
                     ( \ hvs -> let ce_out = extendClosureEnv ce_in (zipLazy nms hvs)
                                in  mapM (linkBCO ie ce_out) ul_bcos )
-
         let ce_all_additions = zip nms hvals
             ce_top_additions = filter (isExternalName.fst) ce_all_additions
             ce_additions     = if toplevs_only then ce_top_additions 
@@ -650,8 +720,22 @@ linkSomeBCOs toplevs_only ie ce_in ul_bcos
 		     -- closure environment, which leads to trouble.
 		     ASSERT (all (not . (`elemNameEnv` ce_in)) (map fst ce_additions))
 		     extendClosureEnv ce_in ce_additions
-        return (ce_out, hvals)
-
+            refs  = goForRefs ul_bcos
+            names = nub$ concatMap (ssElts . unlinkedBCOItbls) (ul_bcos ++ refs)
+        addresses <- mapM (lookupIE ie) names
+        let de_additions = [(address, name) | (address, name) <- zip addresses names
+                                            , not(address `elemAddressEnv` de_in) 
+                           ]
+            de_out = extendAddressEnvList' de_in de_additions
+        return ( ce_out, de_out, hvals)
+    where 
+          goForRefs = getRefs []
+          getRefs acc []  = acc
+          getRefs acc new = getRefs (new++acc) 
+                 [bco | BCOPtrBCO bco <- concatMap (ssElts . unlinkedBCOPtrs) new
+                      , notElemBy bco (new ++ acc) nameEq]
+          ul1 `nameEq` ul2 = unlinkedBCOName ul1 == unlinkedBCOName ul2
+          (x1 `notElemBy` x2) eq = null$ intersectBy eq [x1] x2
 \end{code}
 
 

compiler/ghci/ObjLink.lhs

@@ -18,9 +18,11 @@ module ObjLink (
    unloadObj,   	 -- :: String -> IO ()
    insertSymbol,         -- :: String -> String -> Ptr a -> IO ()
    lookupSymbol,	 -- :: String -> IO (Maybe (Ptr a))
-   resolveObjs  	 -- :: IO SuccessFlag
+   resolveObjs,  	 -- :: IO SuccessFlag
+   lookupDataCon         -- :: Ptr a  -> IO (Maybe String)
   )  where
 
+import ByteCodeItbls    ( StgInfoTable )
 import Panic		( panic )
 import BasicTypes	( SuccessFlag, successIf )
 import Config		( cLeadingUnderscore )
@@ -31,6 +33,10 @@ import Foreign.C
 import Foreign		( nullPtr )
 import GHC.Exts         ( Ptr(..), unsafeCoerce# )
 
+import Constants        ( wORD_SIZE )
+import Foreign          ( plusPtr )
+
+
 -- ---------------------------------------------------------------------------
 -- RTS Linker Interface
 -- ---------------------------------------------------------------------------
@@ -51,6 +57,14 @@ lookupSymbol str_in = do
 	then return Nothing
 	else return (Just addr)
 
+-- | Expects a Ptr to an info table, not to a closure
+lookupDataCon :: Ptr StgInfoTable -> IO (Maybe String)
+lookupDataCon ptr = do
+    name <- c_lookupDataCon  (ptr `plusPtr` (wORD_SIZE*2))
+    if name == nullPtr
+       then return Nothing
+       else peekCString name >>= return . Just
+
 prefixUnderscore :: String -> String
 prefixUnderscore
  | cLeadingUnderscore == "YES" = ('_':)
@@ -94,5 +108,6 @@ foreign import ccall unsafe "lookupSymbol" c_lookupSymbol :: CString -> IO (Ptr
 foreign import ccall unsafe "loadObj"      c_loadObj :: CString -> IO Int
 foreign import ccall unsafe "unloadObj"    c_unloadObj :: CString -> IO Int
 foreign import ccall unsafe "resolveObjs"  c_resolveObjs :: IO Int
+foreign import ccall unsafe "lookupDataCon"  c_lookupDataCon :: Ptr a -> IO CString
 
 \end{code}

compiler/prelude/TysWiredIn.lhs

@@ -38,6 +38,8 @@ module TysWiredIn (
 	unitTyCon, unitDataCon, unitDataConId, pairTyCon, 
 	unboxedSingletonTyCon, unboxedSingletonDataCon,
 	unboxedPairTyCon, unboxedPairDataCon,
+       
+        boxedTupleArr, unboxedTupleArr,
 
 	unitTy,
 

includes/Linker.h

@@ -33,6 +33,9 @@ HsInt resolveObjs( void );
 /* load a dynamic library */
 char *addDLL( char* dll_name );
 
+/* lookup an address in the datacon tbl */
+char *lookupDataCon( StgWord addr);
+
 extern void markRootPtrTable(void (*)(StgClosure **));
 
 #endif /* LINKER_H */