Variables.hs

{-# LANGUAGE LambdaCase        #-}
{-# LANGUAGE MagicHash         #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards   #-}
{-# LANGUAGE TemplateHaskell   #-}
{-# OPTIONS -fno-warn-name-shadowing #-}
{-# OPTIONS -fno-warn-unused-matches #-}
{-# OPTIONS -fno-warn-unused-imports #-}
    -- Any moved from GHC.Prim to GHC.Types
    -- so import both and use unused imports to get compatibility

-- | Module for debugging Haskell programs. To use, take the functions that
--   you are interested in debugging, e.g.:
--
-- > module QuickSort(quicksort) where
-- > import Data.List
-- >
-- > quicksort :: Ord a => [a] -> [a]
-- > quicksort [] = []
-- > quicksort (x:xs) = quicksort lt ++ [x] ++ quicksort gt
-- >     where (lt, gt) = partition (<= x) xs
--
--   Turn on the @TemplateHaskell@ and @ViewPatterns@ extensions, import "Debug",
--   indent your code and place it under a call to 'debug', e.g.:
--
-- > {-# LANGUAGE TemplateHaskell, ViewPatterns, PartialTypeSignatures #-}
-- > {-# OPTIONS_GHC -Wno-partial-type-signatures #-}
-- > module QuickSort(quicksort) where
-- > import Data.List
-- > import Debug
-- >
-- > debug [d|
-- >    quicksort :: Ord a => [a] -> [a]
-- >    quicksort [] = []
-- >    quicksort (x:xs) = quicksort lt ++ [x] ++ quicksort gt
-- >        where (lt, gt) = partition (<= x) xs
-- >    |]
--
--   We can now run our debugger with:
--
-- > $ ghci QuickSort.hs
-- > GHCi, version 8.2.1: http://www.haskell.org/ghc/  :? for help
-- > [1 of 1] Compiling QuickSort        ( QuickSort.hs, interpreted )
-- > Ok, 1 module loaded.
-- > *QuickSort> quicksort "haskell"
-- > "aehklls"
-- > *QuickSort> debugView
--
--   The final call to 'debugView' starts a web browser to view the recorded information.
--   Alternatively call 'debugSave' to write the web page to a known location.
--
--   For more ways to view the result (e.g. producing JSON) or record traces (without using
--   @TemplateHaskell@) see "Debug.DebugTrace".
module Debug.Variables(
    debug,
    debugClear,
    debugRun,
    debugPrint,
    debugJSON,
    debugView,
    debugSave,
    DebugTrace(..),
    getDebugTrace,
    -- * Recording
    funInfo, fun, var,
    ) where

import           Control.DeepSeq
import           Control.Exception
import           Control.Monad
import           Control.Monad.Extra
import           Data.Aeson
import           Data.Aeson.Text
import           Data.Aeson.Types
import qualified Data.ByteString.Lazy.Char8 as B
import           Data.Char
import Data.Generics.Uniplate.Data
import           Data.Hashable
import qualified Data.HashMap.Strict        as HM
import           Data.IORef
import           Data.List.Extra
import           Data.Maybe
import           Data.Monoid                ()
import           Data.Text                  (Text, pack)
import qualified Data.Text                  as T
import qualified Data.Text.Lazy             as TL
import qualified Data.Text.Lazy.IO          as TL
import           Data.Text.Read             as T
import           Data.Tuple.Extra
import qualified Data.Vector                as V
import           Debug.DebugTrace
import           Debug.Util
import           GHC.Generics
import           GHC.Prim
import           GHC.Types
import           Language.Haskell.TH
import           Language.Haskell.TH.Syntax
import           System.Directory
import           System.IO
import           System.IO.Unsafe
import           Text.Show.Functions        ()
import           Unsafe.Coerce

-- | Run a computation and open a browser window showing observed function calls.
--
--   @ main = debugRun $ do
--       ...
--   @
debugRun :: IO a -> IO a
debugRun = bracket_ debugClear debugView

-- | Print information about the observed function calls to 'stdout',
--   in a human-readable format.
debugPrint :: IO ()
debugPrint = getDebugTrace >>= debugPrintTrace

-- | Save information about observed functions to the specified file, in HTML format.
debugSave :: FilePath -> IO ()
debugSave fp = debugSaveTrace fp =<< getDebugTrace

-- | Open a web browser showing information about observed functions.
debugView :: IO ()
debugView = getDebugTrace >>= debugViewTrace

-- | Obtain information about observed functions in JSON format.
--   The JSON format is not considered a stable part of the interface,
--   more presented as a back door to allow exploration of alternative
--   views.
debugJSON :: IO String
debugJSON = B.unpack . debugJSONTrace <$> getDebugTrace

------------------------------------------------------------------
-- DebugTrace Backend

-- | A single function call, used to attach additional information
data Call = Call Function (IORef [(Text, Var)])

{-# NOINLINE refVariables #-}
refVariables :: IORef Variables
refVariables = unsafePerformIO $ newIORef newVariables

{-# NOINLINE refCalls #-}
refCalls :: IORef [Call]
refCalls = unsafePerformIO $ newIORef []

-- | Clear all debug information. Useful when working in @ghci@ to reset
--   any previous debugging work and reduce the amount of output.
debugClear :: IO ()
debugClear = do
    writeIORef refVariables newVariables
    writeIORef refCalls []

-- | Returns all the information about the observed function accumulated so far
--   in the variables.
getDebugTrace :: IO DebugTrace
getDebugTrace = do
  vars <- readIORef refVariables
  vars <- return $ map varShow $ listVariables vars
  calls <- readIORef refCalls
  let infos = nubOrd [x | Call x _ <- calls]
      infoId = HM.fromList $ zip infos [0::Int ..]
  callEntries <-
    forM (reverse calls) $ \(Call info vars) -> do
      vars <- readIORef vars
      let callFunctionId   = infoId HM.! info
          callVals = map (second varId) vars
          callDepends = [] -- available in the Hoed backend but not in this one
          callParents = [] -- available in the Hoed backend but not in this one
      return CallData{..}
  return $ DebugTrace infos (map T.pack vars) callEntries

------------------------------------------------------------------
-- Instrumentation

{-# NOINLINE fun #-}
-- | Called under a lambda with a function name to provide a unique context for
--   a particular call, e.g.:
--
-- > tracedAdd x y = fun "add" $ \t -> var t "x" x + var t "y" y
--
--   This function involves giving identity to function calls, so is unsafe,
--   and will only work under a lambda.
fun :: Show a => String -> (Call -> a) -> a
fun name = funInfo $ Function (T.pack name) "" [] ""

-- | A version of 'fun' allowing you to pass further information about the
--   'Function' which is used when showing debug views.
funInfo :: Show a => Function -> (Call -> a) -> a
{-# NOINLINE funInfo #-}
funInfo info f = unsafePerformIO $ do
    ref <- newIORef []
    let x = Call info ref
    atomicModifyIORef refCalls $ \v -> (x:v, ())
    return $ f x

{-# NOINLINE var #-}
-- | Used in conjunction with 'fun' to annotate variables. See 'fun' for an example.
var :: Show a => Call -> String -> a -> a
var (Call _ ref) name val = unsafePerformIO $ do
    when (show val /= "<function>") $ do -- these make the variable list long without providing useful info
        var <- atomicModifyIORef refVariables $ addVariable val
        name' <- unShadowName ref $ pack name
        whenJust name' (\n -> atomicModifyIORef ref $ \v -> ((n, var) :v, ()))
    return val

-- | If a name is already used, find the next available name by adding ' (prime) chars until
--   the resulting name is unique
unShadowName :: IORef [(Text, Var)] -> Text -> IO (Maybe Text)
unShadowName ioRef t = do
    pairs <- readIORef ioRef
    let matches = filter (isPrefixPrime t) $ map fst pairs
    let shadowedLimit = 3
    if not (null matches)
        then do
            let lengths = map T.length matches
            let zipped = zip matches lengths
            let maxLen = maximum lengths
            let maxName = fst $ fromJust $ find (\p -> snd p == maxLen) zipped
            case length matches of    -- --i.e., case (the number of (')s)
                n | n > shadowedLimit  -> return Nothing
                _                      -> return $ Just $ maxName `T.append` "'"
        else return $ Just t

-- | Is the second string equal to the first plus some number of ' (prime) characters?
-- | e.g. x `isPrefixPrime` x' == true, x isPrefixPrime x'' == True, but x isPrefixPrime xs == False
isPrefixPrime :: Text -> Text -> Bool
isPrefixPrime s t = s == T.dropWhileEnd (== '\'') t

-- | A @TemplateHaskell@ wrapper to convert a normal function into a traced function.
--   For an example see "Debug". Inserts 'funInfo' and 'var' calls.
debug :: Q [Dec] -> Q [Dec]
debug q = do
    missing <- filterM (notM . isExtEnabled) [ViewPatterns, PartialTypeSignatures]
    when (missing /= []) $
        error $ "\ndebug [d| ... |] requires additional extensions:\n" ++
                "{-# LANGUAGE " ++ intercalate ", " (map show missing) ++ " #-}\n"
    decs <- q
    let askSig x = find (\case SigD y _ -> x == y; _ -> False) decs
    mapM (adjustDec askSig) decs

adjustDec :: (Name -> Maybe Dec) -> Dec -> Q Dec
-- try and shove in a "_ =>" if we can, to capture necessary Show instances
adjustDec askSig x@(SigD name ty@(ForallT vars ctxt typ))
  | hasRankNTypes ty = return x
  | otherwise = return $
    SigD name $ ForallT vars (delete WildCardT ctxt ++ [WildCardT]) typ
adjustDec askSig (SigD name typ) = adjustDec askSig $ SigD name $ ForallT [] [] typ
adjustDec askSig o@(FunD name clauses@(Clause arity _ _:_))
  | Just (SigD _ ty) <- askSig name
  , hasRankNTypes ty = return o
  | otherwise = do
    inner <- newName "inner"
    tag <- newName "tag"
    args <- sequence [newName $ "arg" ++ show i | i <- [1 .. length arity]]
    let addTag (Clause ps bod inner) = Clause (VarP tag:ps) bod inner
    let clauses2 = map addTag $ transformBi (adjustPat tag) clauses
    let args2 = [VarE 'var `AppE` VarE tag `AppE` toLitPre "$" a `AppE` VarE a | a <- args]
    let info = ConE 'Function `AppE`
            packLit (toLit name) `AppE`
            packLit (LitE (StringL $ prettyPrint $ maybeToList (askSig name) ++ [o])) `AppE`
            ListE (map (packLit . toLitPre "$") args) `AppE`
            packLit (LitE (StringL "$result"))
    let body2 = VarE 'var `AppE` VarE tag `AppE` LitE (StringL "$result") `AppE` foldl AppE (VarE inner) (VarE tag : args2)
    let body = VarE 'funInfo `AppE` info `AppE` LamE [VarP tag] body2
    afterApps <- transformApps tag clauses2
    return $ FunD name [Clause (map VarP args) (NormalB body) [FunD inner afterApps]]
adjustDec askSig x = return x

transformApps :: Name -> [Clause] -> Q [Clause]
transformApps tag = mapM (appsFromClause tag)

appsFromClause :: Name -> Clause -> Q Clause
appsFromClause tag cl@(Clause pats body decs) = do
    newBody <- appsFromBody tag body
    newDecs <- mapM (appsFromDec tag) decs
    return $ Clause pats newBody newDecs

appsFromBody :: Name -> Body -> Q Body
appsFromBody _ b@(GuardedB _) = return b -- TODO: implement guards
appsFromBody tag (NormalB e)  = NormalB <$> appsFromExp tag e

appsFromExp :: Name -> Exp -> Q Exp
appsFromExp tag e@(AppE e1 e2) = do
    newE1 <- appsFromExp tag e1
    newE2 <- appsFromExp tag e2
    adjustApp tag (AppE newE1 newE2)
appsFromExp tag e@(LetE decs exp) = do
    newDecs <- traverse (appsFromDec tag) decs
    LetE newDecs <$> appsFromExp tag exp
appsFromExp tag e@(InfixE e1May e2 e3May) = do
    newE1 <- appsFromExpMay tag e1May
    newE2 <- appsFromExp tag e2
    newE3 <- appsFromExpMay tag e3May
    adjustApp tag (InfixE newE1 newE2 newE3)
appsFromExp tag e@(CaseE exp matches) = do
    newExp <- appsFromExp tag exp
    newMatches <- traverse (appsFromMatch tag) matches
    return $ CaseE newExp newMatches
appsFromExp tag e = return e

appsFromExpMay :: Name -> Maybe Exp -> Q (Maybe Exp)
appsFromExpMay tag Nothing  = return Nothing
appsFromExpMay tag (Just e) = sequence $ Just $ appsFromExp tag e

appsFromDec :: Name -> Dec -> Q Dec
appsFromDec tag d@(ValD pat body dec) = do
    newBody <- appsFromBody tag body
    return $ ValD pat newBody dec
appsFromDec tag d@(FunD name subClauses) =
    FunD name <$> traverse (appsFromClause tag) subClauses
appsFromDec _ d = return d

appsFromMatch :: Name -> Match -> Q Match
appsFromMatch tag (Match pat body decs) = do
    newBody <- appsFromBody tag body
    newDecs <- traverse (appsFromDec tag) decs
    return $ Match pat newBody newDecs

adjustApp :: Name -> Exp -> Q Exp
adjustApp tag (AppE e1 e2) = do
    let displayName = expDisplayName e1
    e1n <- newName displayName
    let viewP = ViewP (VarE 'var `AppE` VarE tag `AppE` LitE (StringL displayName)) (VarP e1n)
    let result = LetE [ValD viewP (NormalB (AppE e1 e2)) []] (VarE e1n)
    return result
adjustApp tag e@(InfixE e1May e2 e3May) = do
    let displayName = infixExpDisplayName e2 -- infix symbol, e.g "++"
    if displayName == "$" --don't record $ as a function application
        then return e
        else do
            let legalInfixVar = mkLegalInfixVar displayName -- infix name as valid variable, e.g. "plus_plus"
            e2Var <- newName legalInfixVar
            let viewP = ViewP (VarE 'var `AppE` VarE tag `AppE` LitE (StringL displayName)) (VarP e2Var)
            return $ LetE [ValD viewP (NormalB (InfixE e1May e2 e3May)) []] (VarE e2Var)
adjustApp _ e@(UInfixE e1 e2 e3) = return e   --TODO: These might need to be processed
adjustApp _ e = return e

-- Find the (unqualified) function name to use as the UI display name
expDisplayName :: Exp -> String
expDisplayName e =
    let name = removeLet $ (show . ppr) e
    in removeExtraDigits (takeWhileEnd (/= '.') ((head . words) name))

-- Same as expDisplayName but for infix functions
infixExpDisplayName :: Exp -> String
infixExpDisplayName e =
    let name = removeLet $ (show . ppr) e
        name' = removeExtraDigits (takeWhileEnd (/= '.') ((head . words) name))
    in fromMaybe name' $ stripSuffix ")" name'

adjustPat :: Name -> Pat -> Pat
adjustPat tag (VarP x) = ViewP (VarE 'var `AppE` VarE tag `AppE` toLit x) (VarP x)
adjustPat tag x = x

toLit :: Name -> Exp
toLit = toLitPre ""
toLitPre :: String -> Name -> Exp
toLitPre pre (Name (OccName x) _) = LitE $ StringL $ pre ++ x
packLit :: Exp -> Exp
packLit = AppE (VarE 'pack)

------------------------------------------------------------------
-- Data structures

data Variables = Variables
    Int -- Number in the array
    [(Any, String)] -- Entries, (key a, show a), indexed from [n..0]

data Var = Var Int String -- index into Variables, show a
         deriving (Eq,Ord)

instance Show Var where
    show (Var i s) = s ++ " @" ++ show i

varId :: Var -> Int
varId (Var x _) = x

varShow :: Var -> String
varShow (Var _ x) = x

newVariables :: Variables
newVariables = Variables 0 []

listVariables :: Variables -> [Var]
listVariables (Variables n xs) = [Var i s | (i,(_,s)) <- zipFrom 0 $ reverse xs]

addVariable :: Show a => a -> Variables -> (Variables, Var)
addVariable a vs@(Variables n xs) =
    case findIndex (\(key,_) -> ptrEqual key keyA) xs of
        Nothing -> (Variables (n+1) ((keyA,showA):xs), Var n showA)
        Just i  -> (vs, Var (n-i-1) $ snd $ xs !! i)
    where
        keyA = unsafeCoerce a
        showA = show a

ptrEqual :: Any -> Any -> Bool
ptrEqual a b = unsafePerformIO $ do
    a <- evaluate a
    b <- evaluate b
    return $ isTrue# (reallyUnsafePtrEquality# a b)