Improve Trie manipulation in Haskell code

Question

I have the following data structure:

data XComposeString = XComposeString { keyEvents :: [Event], s :: Result }
  deriving (Eq, Show)

data XComposeFile = XComposeFile { strings :: [XComposeString] }
  deriving (Eq, Show)

Where Event and Result just String or some combination of several String.

I'm creating a Trie from list of type [(keyEvents x, [s x])]:

constructTrie list = M.fromListWith comp list
  where
    comp = \old new -> old ++ new

I want to check on duplicates and prefix overlaps.

duplicates :: M.TrieMap Map.Map Event [Result] -> String
duplicates m = M.showTrie list ""
  where
    list = M.filter (\v -> length v /= 1) m

prefixOverlap :: M.TrieMap Map.Map Event [Result] -> String
prefixOverlap m = M.showTrie list ""
  where
    list = M.filterWithKey (\k v -> not $ null $ catMaybes [M.lookup x m | x <- reverse (inits (init k))]) m

How can I improve this, so, there would be three functions: one to create, one to check on duplicates and one to check on prefix overlaps.

Useful imports and types definition:

import Data.List(inits)
import Data.Maybe(catMaybes)
import qualified Data.Map as Map
import qualified Data.ListTrie.Map as M

type Event = String
type Result = String

Also, cabal install list-tries is needed for ListTrie data structure.

Answer 1

Can you add minimal imports so your code compiles? I had to guess:

import Data.ListTrie.Map as M hiding (null)
import qualified Data.Map as Map
import Data.Maybe
import Data.List


data Event = Event deriving (Show, Eq, Ord)
data Result = Result deriving (Show, Eq)

constructTrie can be easily improved by inlining comp:

constructTrie list = M.fromListWith (++) list

Or even to constructTrie = M.fromListWith (++) but that requires turning off monomorphism restriction or adding an explicit signature for constructTrie, so it's up to your taste.

Also you can deduplicate showTrie calls and long trie signatures:

mshow list = M.showTrie list ""

type ResultTrie = M.TrieMap Map.Map Event [Result]

duplicates :: ResultTrie -> String
duplicates m = mshow list
  where
    list = M.filter (\v -> length v /= 1) m

prefixOverlap :: ResultTrie -> String
prefixOverlap m = mshow list
  where
    list = M.filterWithKey (\k v -> not $ null $ catMaybes [M.lookup x m | x <- reverse (inits (init k))]) m

Now it's beneficial to inline list variables in both functions:

duplicates :: ResultTrie -> String
duplicates m = mshow $ M.filter (\v -> length v /= 1) m

prefixOverlap :: ResultTrie -> String
prefixOverlap m = mshow $ M.filterWithKey (\k v -> not $ null $ catMaybes [M.lookup x m | x <- reverse (inits (init k))]) m

Now you can get rid of m in duplicates:

duplicates :: ResultTrie -> String
duplicates = mshow . M.filter (\v -> length v /= 1)

In prefixOverlap the line is too long, I can split it. I chose to extract the (\ -> ) construct. Note that m is in more that 1 place so we cannot use ..

prefixOverlap :: ResultTrie -> String
prefixOverlap m = mshow $ M.filterWithKey f m where
    f k _ = not $ null $ catMaybes [M.lookup x m | x <- reverse (inits (init k))]

Now remove extra parentheses in the comprehension:

prefixOverlap :: ResultTrie -> String
prefixOverlap m = mshow $ M.filterWithKey f m where
    f k _ = not $ null $ catMaybes [M.lookup x m | x <- reverse $ inits $ init k]

From here I'll only work on f.

Now a trickier part. Long chains of $ are equivalent to . following one $ in the end:

    f k _ = (not . null . catMaybes) [M.lookup x m | x <- reverse $ inits $ init k]

not . null . catMaybes is a function of [Maybe a] -> Bool that returns True only if there is any Just in the list. I can literally translate this sentence into Haskell: it's just any isJust:

    f k _ = any isJust [M.lookup x m | x <- reverse $ inits $ init k]

Now even trickier: your comprehension is just a map:

    f k _ = any isJust $ map (\x -> M.lookup x m) $ reverse $ inits $ init k

Congratulations, we can again turn a chain of $ into a chain of .:

    f k _ = any isJust . map (\x -> M.lookup x m) . reverse . inits . init $ k

Now remember that . is associative so you can put parentheses in the middle of the chain wherever you want:

any isJust . (map (\x -> M.lookup x m) . reverse) . inits . init

Now we have map something . reverse. And it doesn't matter when you reverse - before or after map:

map something f . reverse = reverse . map something

That is, map and reverse commute and we can swap them and remove parens:

any isJust . reverse . map (\x -> M.lookup x m) . inits . init

Now we can put around any and reverse. It's clear that any works the same way for straight and reversed lists, that is, any something . reverse = any something, so we can throw reverse away:

any isJust . map (\x -> M.lookup x m) . inits . init

And put back into prefixOverlap:

prefixOverlap :: ResultTrie -> String
prefixOverlap m = mshow $ M.filterWithKey f m where
    f k _ = any isJust . map (\x -> M.lookup x m) . inits . init $ k

Here is the full version of the code:

import Data.ListTrie.Map as M hiding (null, map)
import qualified Data.Map as Map
import Data.Maybe
import Data.List

type Event = Char
type Result = ()
type ResultTrie = M.TrieMap Map.Map Event [Result]

data XComposeString = XComposeString { keyEvents :: [Event], s :: Result }
  deriving (Eq, Show)

data XComposeFile = XComposeFile { strings :: [XComposeString] }
  deriving (Eq, Show)

constructTrie list = M.fromListWith (++) list

mshow :: ResultTrie -> String
mshow list = M.showTrie list ""


duplicates :: ResultTrie -> String
duplicates = mshow . M.filter (\v -> length v /= 1)

prefixOverlap :: ResultTrie -> String
prefixOverlap m = mshow $ M.filterWithKey f m where
    f k _ = any isJust . map (\x -> M.lookup x m) . inits . init $ k

Note that I used fake Event and Result types. I also checked that old and new prefixOverlap are the same using QuickCheck and smallcheck libraries:

prefixOverlapOld :: ResultTrie -> String
prefixOverlapOld m = M.showTrie list ""
  where
    list = M.filterWithKey (\k v -> not $ null $ catMaybes [M.lookup x m | x <- reverse (inits (init k))]) m

prop_foo x = prefixOverlap m == prefixOverlapOld m where
    m = M.fromList $ map (first (++ "X")) x

first is from Control.Arrow.

The prefixOverlap probably can still be improved with using lookupPrefix and/or children functions from Data.ListTrie.Map.

Here is my take:

prefixOverlapNew :: ResultTrie -> String
prefixOverlapNew = mshow . M.unions . map (\(p,m) -> M.addPrefix [p] m) . Prelude.filter (\(p, m) -> M.size m /= 1 || (null $ fst $ head $ M.toList m)) . Map.toList . M.children1

Note that it differs from both original prefixOverlapOld and refactored prefixOverlap:

QuickCheck found that they differ on a map with 2 keys "fo" and "f". The prefixOverlapNew version correctly shows both "fo" and "f" as they overlap. The original version(s) show only value on "fo" key which is dubuios.

Here is an improved version using list comprehension to combine filter and map:

prefixOverlapNew :: ResultTrie -> String
prefixOvelapNew mm = mshow $ 
    M.unions [M.addPrefix [p] m | (p, m) <- Map.toList $ M.children1 mm,  M.size m /= 1 || null (fst $ head $ M.toList m)]

Moved f out:

prefixOverlapNew mm = mshow $ M.unions [M.addPrefix [p] m | (p, m) <- Map.toList $ M.children1 mm, f m] where
        f m =  M.size m /= 1 || null (fst $ head $ M.toList m)

Finally, a version in list monad, just for fun:

prefixOverlapNew :: ResultTrie -> String
prefixOverlapNew mm = mshow $ M.unions $ do
        (p, m) <- Map.toList $ M.children1 mm
        guard $ M.size m /= 1 || null (fst $ head $ M.toList m)
        return $ M.addPrefix [p] m

guard is from Control.Monad.