Prettify.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE ScopedTypeVariables #-}

-- | Functions to make environments easier to read
module Language.Fixpoint.Solver.Prettify (savePrettifiedQuery) where

import           Data.Bifunctor (first)
import           Data.HashMap.Lazy (HashMap)
import qualified Data.HashMap.Lazy as HashMap
import           Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
import           Data.List (intersperse, sortOn)
import           Data.Maybe (fromMaybe)
import           Data.Text (Text)
import qualified Data.Text as Text
import           Language.Fixpoint.Misc (ensurePath)
import           Language.Fixpoint.Solver.EnvironmentReduction
  ( dropLikelyIrrelevantBindings
  , inlineInSortedReft
  , mergeDuplicatedBindings
  , simplifyBooleanRefts
  , undoANFAndVV
  )
import           Language.Fixpoint.Types.Config (Config, queryFile)
import           Language.Fixpoint.Types.Constraints
import           Language.Fixpoint.Types.Environments
  (BindEnv, elemsIBindEnv, lookupBindEnv)
import           Language.Fixpoint.Types.Names
  ( Symbol
  , dropPrefixOfSym
  , prefixOfSym
  , suffixOfSym
  , suffixSymbol
  , symbol
  , symbolText
  )
import           Language.Fixpoint.Types.PrettyPrint
import           Language.Fixpoint.Types.Refinements
  ( Expr(..)
  , pattern PFalse
  , Reft
  , SortedReft(..)
  , conjuncts
  , expr
  , reft
  , reftBind
  , reftPred
  , sortedReftSymbols
  , substf
  , substSortInExpr
  )
import           Language.Fixpoint.Types.Sorts (Sort(..), substSort)
import           Language.Fixpoint.Types.Substitutions (exprSymbolsSet)
import qualified Language.Fixpoint.Utils.Files as Files
import           System.FilePath (addExtension)
import           Text.PrettyPrint.HughesPJ.Compat
  (Doc, braces, hang, nest, render, sep, text, vcat, (<+>), ($+$))


savePrettifiedQuery :: Fixpoint a => Config -> FInfo a -> IO ()
savePrettifiedQuery cfg info = do
  let fq   = queryFile Files.Fq cfg `addExtension` "prettified"
  putStrLn $ "Saving prettified Query: "   ++ fq ++ "\n"
  ensurePath fq
  writeFile fq $ render (prettyConstraints info)

prettyConstraints :: Fixpoint a => FInfo a -> Doc
prettyConstraints info =
  vcat $
  map
    (prettyConstraint (bs info) . snd)
    (sortOn fst $ HashMap.toList (cm info))

prettyConstraint
  :: Fixpoint a
  => BindEnv a
  -> SubC a
  -> Doc
prettyConstraint bindEnv c =
  let env = [ (s, ([(bId, a)], sr))
            | bId <- elemsIBindEnv $ senv c
            , let (s, sr, a) = lookupBindEnv bId bindEnv
            ]
      mergedEnv = mergeDuplicatedBindings env
      undoANFEnv = undoANFAndVV mergedEnv
      boolSimplEnvDiff = simplifyBooleanRefts undoANFEnv
      boolSimplEnv = HashMap.map snd $ HashMap.union boolSimplEnvDiff undoANFEnv

      simplifiedLhs = simplify $ inlineInSortedReft (`HashMap.lookup` boolSimplEnv) (slhs c)
      simplifiedRhs = simplify $ inlineInSortedReft (`HashMap.lookup` boolSimplEnv) (srhs c)

      prunedEnv =
        if expr simplifiedRhs /= PFalse then
          dropLikelyIrrelevantBindings
            (constraintSymbols simplifiedLhs simplifiedRhs)
            boolSimplEnv
        else
          boolSimplEnv
      (renamedEnv, c') =
        shortenVarNames prunedEnv c { slhs = simplifiedLhs, srhs = simplifiedRhs }
      prettyEnv =
        sortOn bindingSelector $
        eraseUnusedBindings (constraintSymbols (slhs c') (srhs c')) renamedEnv
   in hang (text "\n\nconstraint:") 2 $
          text "lhs" <+> toFix (slhs c')
      $+$ text "rhs" <+> toFix (srhs c')
      $+$ pprId (sid c) <+> text "tag" <+> toFix (stag c)
      $+$ toFixMeta (text "constraint" <+> pprId (sid c)) (toFix (sinfo c))
      $+$ hang (text "environment:") 2
            (vcat $ intersperse "" $ elidedMessage : map prettyBind prettyEnv)
  where
    prettyBind (s, sr) = sep [toFix s <+> ":", nest 2 $ prettySortedRef sr]
    prettySortedRef sr = braces $ sep
      [ toFix (reftBind (sr_reft sr)) <+> ":" <+> toFix (sr_sort sr) <+> "|"
      , nest 2 $ toFix $ conjuncts $ reftPred $ sr_reft sr
      ]

    elidedMessage = "// elided some likely irrelevant bindings"

    constraintSymbols sr0 sr1 =
      HashSet.union (sortedReftSymbols sr0) (sortedReftSymbols sr1)

    -- Sort by symbol then reft
    bindingSelector (s, sr) =
      ( -- put unused symbols at the end
        if "_" `Text.isPrefixOf` s then "{" <> s else s
      , reftBind (sr_reft sr)
      , sr_sort sr
      , reftPred $ sr_reft sr
      )

pprId :: Show a => Maybe a -> Doc
pprId (Just i)  = "id" <+> text (show i)
pprId _         = ""

toFixMeta :: Doc -> Doc -> Doc
toFixMeta k v = text "// META" <+> k <+> text ":" <+> v

-- | @eraseUnusedBindings ss env@ prefixes with @_@ the symbols that aren't
-- refered from @ss@ or refinement types in the environment.
eraseUnusedBindings
  :: HashSet Symbol -> [(Symbol, SortedReft)] -> [(Text, SortedReft)]
eraseUnusedBindings ss env = map (first reachable) env
  where
    allSymbols = HashSet.union ss envSymbols
    envSymbols =
      HashSet.unions $
      map (exprSymbolsSet . reftPred . sr_reft . snd) env

    reachable s =
      if s `HashSet.member` allSymbols then
        symbolText s
      else
        "_" <> symbolText s

-- | Shortens the names of refinements types in a given environment
-- and constraint
shortenVarNames
  :: HashMap Symbol SortedReft
  -> SubC a
  -> ([(Symbol, SortedReft)], SubC a)
shortenVarNames env subc =
  let bindsRenameMap = proposeRenamings $ HashMap.keys env
      env' = map (renameBind bindsRenameMap) (HashMap.toList env)
   in
      (env', renameSubC bindsRenameMap subc)
  where
    renameSubC :: HashMap Symbol Symbol -> SubC a -> SubC a
    renameSubC symMap c =
      mkSubC
        (senv c)
        (renameSortedReft symMap $ slhs c)
        (renameSortedReft symMap $ srhs c)
        (sid c)
        (stag c)
        (sinfo c)

    renameBind
      :: HashMap Symbol Symbol -> (Symbol, SortedReft) -> (Symbol, SortedReft)
    renameBind symMap (s, sr) =
      (at symMap s, renameSortedReft symMap sr)

    renameSortedReft
      :: HashMap Symbol Symbol -> SortedReft -> SortedReft
    renameSortedReft symMap (RR t r) =
      let sortSubst = FObj . at symMap
       in RR (substSort sortSubst t) (renameReft symMap r)

    renameReft :: HashMap Symbol Symbol -> Reft -> Reft
    renameReft symMap r =
      let m = HashMap.insert (reftBind r) (prefixOfSym $ reftBind r) symMap
          sortSubst = FObj . at symMap
       in reft (at m (reftBind r)) $
            substSortInExpr sortSubst $
            substf (EVar . at m) (reftPred r)

    at :: HashMap Symbol Symbol -> Symbol -> Symbol
    at m k = fromMaybe k $ HashMap.lookup k m

-- | Given a list of symbols with no duplicates, it proposes shorter
-- names to use for them.
--
-- All proposals preserve the prefix of the original name. A prefix is
-- the longest substring containing the initial character and no separator
-- ## ('symSepName').
--
-- Suffixes are preserved, except when symbols with a given prefix always
-- use the same suffix.
--
proposeRenamings :: [Symbol] -> HashMap Symbol Symbol
proposeRenamings = toSymMap . toPrefixSuffixMap

-- | Indexes symbols by prefix and then by suffix. If the prefix
-- equals the symbol, then the suffix is the empty symbol.
--
-- For instance,
--
-- > toPrefixSuffixMap ["a##b##c"] ! "a" ! "c" == ["a##b##c"]
-- > toPrefixSuffixMap ["a"] ! "a" ! "" == ["a"]
--
-- In general,
--
-- > forall ss.
-- > Set.fromList ss == Set.fromList $ concat [ xs | m <- elems (toPrefixSuffixMap ss), xs <- elems m ]
--
-- > forall ss.
-- > and [ all (pfx `isPrefixOfSym`) xs && all (sfx `isSuffixOfSym`) xs
-- >     | (pfx, m) <- toList (toPrefixSuffixMap ss)
-- >     , (sfx, xs) <- toList m
-- >     ]
--
-- TODO: put the above in unit tests
--
toPrefixSuffixMap :: [Symbol] -> HashMap Symbol (HashMap Symbol [Symbol])
toPrefixSuffixMap xs = HashMap.fromListWith (HashMap.unionWith (++))
  [ (pfx, HashMap.singleton sfx [s])
  | s <- xs
  , let pfx = prefixOfSym s
        sfx = suffixOfSym $ dropPrefixOfSym s
  ]

-- | Suggests renamings for every symbol in a given prefix/suffix map.
toSymMap :: HashMap Symbol (HashMap Symbol [Symbol]) -> HashMap Symbol Symbol
toSymMap prefixMap = HashMap.fromList
  [ r
  | (pfx, h) <- HashMap.toList prefixMap
  , r <- renameWithSuffixes pfx (HashMap.toList h)
  ]
  where
    renameWithSuffixes pfx = \case
      [(_sfx, ss)] -> renameWithAppendages pfx ("", ss)
      sfxs -> concatMap (renameWithAppendages pfx) sfxs

    renameWithAppendages pfx (sfx, xs) = zip xs $ case xs of
      [_s] -> [pfx `suffixIfNotNull` sfx]
      ss -> zipWith (rename pfx sfx) [1 :: Integer ..] ss

    rename pfx sfx i _s =
      pfx `suffixIfNotNull` sfx `suffixSymbol` symbol (show i)

    suffixIfNotNull a b =
      if Text.null (symbolText b) then a else a `suffixSymbol` b