yeganesh-2.2: yeganesh.hs

-- boilerplate {{{
module Main where

import Control.Arrow ((&&&), second)
import Control.Monad (liftM, filterM, when)
import Data.Char (toLower)
import Data.List (partition, sortBy)
import Data.Map (Map, insert, findWithDefault, fromList, intersection, toList, union)
import Data.Ord (comparing)
import Data.Time (UTCTime, diffUTCTime, getCurrentTime)
import System.Console.GetOpt (ArgDescr(NoArg, ReqArg), ArgOrder(RequireOrder), OptDescr(Option), getOpt, usageInfo)
import System.Directory (createDirectoryIfMissing, doesFileExist, getDirectoryContents, removeDirectory, removeFile)
import System.Environment (getArgs, getEnv)
import System.Environment.XDG.BaseDir (getAllDataFiles, getUserDataDir, getUserDataFile)
import System.Exit (ExitCode(ExitSuccess, ExitFailure), exitWith)
import System.FilePath ((</>))
import System.IO (hClose, hGetContents, hPutStr, stderr, stdout)
import System.Process (runInteractiveProcess, waitForProcess)
import Version (CurrentFormat, parseCurrentFormat, version)
import qualified System.IO.Strict as Strict (getContents, readFile)
-- }}}
-- getopt {{{
options :: [OptDescr Flag]
options =   [Option "p" ["profile"] (ReqArg Profile "PROFILE") "which popularity profile to use"
            ,Option "f" ["filter" ] (NoArg Prune)              "prune a profile to contain exactly the lines of stdin"
            ,Option "v" ["version"] (NoArg Version)            "print the version number"
            ,Option "h" ["help"   ] (NoArg Help)               "show usage information"
            ]

data Options = Options { dmenuOpts :: [String], profile :: String, prune :: Bool }
data Flag    = Profile String | Prune | Version | Help deriving Eq

compactFlags :: [Flag] -> (Flag, Bool)
compactFlags fs = (flag, not $ null prunes) where
    (prunes, nonPrunes)     = partition (==Prune) fs
    flag                    = foldr1 compactFlags' . (Profile "default" :) $ nonPrunes
    compactFlags' Help    _ = Help
    compactFlags' _    Help = Help
    compactFlags' Version _ = Version
    compactFlags' _ Version = Version
    compactFlags' _ p       = p

introText :: String
introText = unlines $ [
    version,
    "Usage: yeganesh [OPTIONS] -- [DMENU_OPTIONS]",
    "OPTIONS are described below, and DMENU_OPTIONS are passed on verbatim to dmenu.",
    "Profiles are stored in the $HOME/.yeganesh directory."]

parseOptions :: [String] -> Either String Options
parseOptions ss = p where
    (opts, dOpts) = fmap (drop 1) . break (== "--") $ ss
    p = case onFirst compactFlags $ getOpt RequireOrder options opts of
        ((Profile f, b), [], []) -> Right (Options dOpts f b)
        ((Version  , _), [], []) -> Left version
        ((Help     , _), [], []) -> Left $ usageInfo introText options
        (_ , ns, []) -> Left $ "Unknown options: " ++ unwords ns
        (_ , _ , es) -> Left  . concat $ es
-- }}}
-- filesystem stuff {{{
type Commands = Map String Double

deprecatedDir :: IO FilePath
deprecatedDir = liftM (</> ".yeganesh") (getEnv "HOME")

inFileName :: String -> IO FilePath
inFileName arg = do
    depDir     <- deprecatedDir
    dataFiles  <- getAllDataFiles "yeganesh" arg
    validFiles <- filterM doesFileExist ((depDir </> arg) : dataFiles)
    case validFiles of
        []    -> getUserDataFile "yeganesh" arg
        (f:_) -> return f

outFileName :: String -> IO FilePath
outFileName arg = do
    dir <- getUserDataDir "yeganesh"
    createDirectoryIfMissing True dir
    return (dir </> arg)

deprecate :: String -> String -> IO ()
deprecate inFile arg = do
    depDir <- deprecatedDir
    when (inFile == depDir </> arg) $ do
        removeFile inFile
        filesLeft <- getDirectoryContents depDir
        when (null . filter (`notElem` [".", ".."]) $ filesLeft)
             (removeDirectory depDir)

readPossiblyNonExistent :: FilePath -> IO CurrentFormat
readPossiblyNonExistent file = catch (Strict.readFile file) (const . return $ "") >>= parseCurrentFormat
-- }}}
-- pure {{{
onFirst :: (a -> a') -> (a, b, c) -> (a', b, c)
onFirst f (a, b, c) = (f a, b, c)

sortOn :: Ord b => (a -> b) -> [a] -> [a]
sortOn f = map snd . sortBy (comparing fst) . map (f &&& id)

descSnd :: Num b => (String, b) -> (b, String)
descSnd = (negate . snd) &&& (map toLower . fst)

showPriority :: Commands -> String
showPriority = unlines . map fst . sortOn descSnd . toList

-- decay exponentially, with a one-month half-life
-- The key for decay is that it be monotonic, so that commands will appear in
-- the same order before and after a decay operation; this means we can delay
-- the decay until *after* the user has selected an option.
decay :: UTCTime -> UTCTime -> Commands -> Commands
decay old new = fmap (/factor) where
    seconds = fromRational . toRational $ diffUTCTime new old
    factor  = exp (seconds * log 2 / 2592000)

-- give a boost, with things close to 0 getting a big boost, and things close
-- to 1 getting a small boost
-- Current method:
--   1. clip ]-infty, infty[ to [0, 1]
--   2. scale [0, 1] to [0.5, 1]
--   3. take sqrt; this is the boost part
--   4. scale [sqrt 0.5, 1] to [0.01, 1]
boost :: (Floating a, Ord a) => a -> a
boost = postscale . sqrt . prescale . clip where
    clip        = min 1 . max 0
    prescale    = (0.5 +) . (/ 2)
    postscale   = ((0.01 - s2) / ms2 +) . ((0.99 / ms2) *)
    s2          = sqrt 0.5
    ms2         = 1 - s2

updatePriority :: String -> UTCTime -> UTCTime -> Commands -> Commands
updatePriority cmd old new cmds = insert cmd pri cmds' where
    cmds' = decay old new cmds
    pri   = boost $ findWithDefault 0 cmd cmds'

parseInput :: String -> Commands
parseInput = fromList . flip zip (repeat 0) . filter (not . null) . lines
-- }}}
-- shell stuff {{{
dmenu :: [String] -> CurrentFormat -> IO (ExitCode, CurrentFormat)
dmenu opts cv@(_, cmds) = do
    (hIn, hOut, hErr, p) <- runInteractiveProcess "dmenu" opts Nothing Nothing
    hPutStr hIn (showPriority cmds)
    hClose hIn
    o <- hGetContents hOut
    e <- hGetContents hErr
    c <- waitForProcess p
    hPutStr stdout o
    hPutStr stderr e
    cv' <- updateState c o cv
    return (c, cv')

updateState :: ExitCode -> String -> CurrentFormat -> IO CurrentFormat
updateState (ExitFailure {})  _  (t, cmds) = return (t, cmds)
updateState  ExitSuccess     cmd (t, cmds) = do
    now <- getCurrentTime
    return (now, updatePriority cmd t now cmds)

runWithOptions :: Options -> IO ()
runWithOptions opts = do
    inFile          <- inFileName  (profile opts)
    outFile         <- outFileName (profile opts)
    cached          <- readPossiblyNonExistent inFile
    new             <- fmap parseInput Strict.getContents
    (code, updated) <- dmenu (dmenuOpts opts) (second (`combine` new) cached)
    writeFile outFile (show updated)
    deprecate inFile (profile opts)
    exitWith code
    where
    combine = if prune opts
              then \old -> union old >>= intersection
              else union
-- }}}
main :: IO ()
main = getArgs >>= either putStrLn runWithOptions . parseOptions