explicitly load primary message part

Previously the primary part was loaded only for text/plain,
because notmuch show included it automatically.

Now any MIME type can be loaded.

1 files changed, 39 insertions, 3 deletions

diff --git a/src/Much/TreeView.hs b/src/Much/TreeView.hs
index 02ccde0..b3d36e8 100644
--- a/src/Much/TreeView.hs
+++ b/src/Much/TreeView.hs
@@ -133,11 +133,11 @@ isQuoteLine s0 = do
 -- Loading / Unloading
 --
 
-
 loadSubForest :: TreeView -> IO (Either String (Forest TreeView))
 loadSubForest = \case
     TVMessage m ->
-        Right
+        pure . Right =<<
+            loadPrimaryMessagePart m =<< id
             . unloadPartsWithFilename
             . map unloadReadSubForests
             . concatMap subForest
@@ -172,7 +172,13 @@ loadSubForest = \case
                     $ xconvPart m mp'
 
     TVSearchResult sr ->
-        Right
+        pure . Right =<<
+            mapM
+              (\case
+                Node tv@(TVMessage m) parts -> Node tv <$> loadPrimaryMessagePart m parts
+                node                        -> pure node
+              )
+            =<< id
             . unloadPartsWithFilename
             . map unloadReadSubForests
             . fromMessageForest
@@ -193,6 +199,36 @@ loadSubForest = \case
     termFromSearchResult = unThreadID . searchThread
 
 
+loadPrimaryMessagePart :: Message -> Forest TreeView -> IO (Forest TreeView)
+loadPrimaryMessagePart m parts =
+    mapM (rewriteTreeIO isPrimaryPart loadNode) parts
+  where
+    pid = partID (primaryMessagePart (messageBody m))
+    mid = messageId m
+
+    isPrimaryPart :: TreeView -> Bool
+    isPrimaryPart = \case
+        TVMessagePart m' mp ->
+            messageId m' == mid && partID mp == pid
+        _ -> False
+
+    loadNode :: Tree TreeView -> IO (Tree TreeView)
+    loadNode node@(Node tv _) =
+        loadSubForest tv >>= \case
+            Right subForest' -> pure (Node tv subForest')
+            Left _           -> pure node
+
+
+rewriteTreeIO :: (a -> Bool) -> (Tree a -> IO (Tree a)) -> Tree a -> IO (Tree a)
+rewriteTreeIO shouldRewrite f t@(Node v children)
+  | shouldRewrite v = do
+      Node v' children' <- f t
+      Node v' <$> mapM (rewriteTreeIO shouldRewrite f) children'
+
+  | otherwise = do
+      Node v <$> mapM (rewriteTreeIO shouldRewrite f) children
+
+
 unloadSubForest :: Tree TreeView -> Forest TreeView
 unloadSubForest t = case rootLabel t of
     TVMessage _ ->