matsuri-0.0.1: Widgets/Tree.hs

{-# LANGUAGE RecordWildCards #-}

-- | Tree.hs
-- Tree widget.
module Widgets.Tree (
    Tree(..),
    NodeData(..),
    defaultTree,
    moveUp,
    moveDn,
    cur
) where

import Prelude hiding (sequence, mapM)
import Control.Monad.State.Lazy hiding (sequence, mapM)
import Data.Traversable (sequence, mapM)
import Data.Tree hiding (Tree)
import Data.Maybe
import Graphics.Vty
import Graphics.Vty.Widgets.Base


-- | Дерево хранится в виде Forest (NodeData a).
-- Поскольку используется готовое дерево, приходится вручную везде
-- проверять инвариант "если дочерних узлов нет, то узел считается
-- развёрнутым".
data Tree
  = Tree
    { selectedAttr   :: Attr
    , nodeAttr       :: Attr
    , selectedIndex  :: Int
    , windowStart    :: Int
    , treeWidth      :: Maybe Int
    , treeHeight     :: Maybe Int
    , nodes          :: Forest NodeData
    }

data NodeData = NodeData {collapsed :: Bool, value :: (String, Value)}
type Value = String

defaultTree tree
  = Tree
    { selectedAttr  = def_attr `with_back_color` cyan `with_fore_color` blue
    , nodeAttr      = def_attr
    , selectedIndex = 0
    , windowStart   = 0
    , treeWidth     = Just 30
    , treeHeight    = Nothing
    , nodes         = tree
    }


----------------------------------------------------------------------
-- moves
-- TODO: first, end, pageup, pagedn

moveUp t@(Tree {windowStart = ws, selectedIndex = sIx})
  | sIx == ws && ws == 0 = t
  | sIx == ws            = t { windowStart = ws-1, selectedIndex = sIx-1 }
  | otherwise            = t { selectedIndex = sIx-1 }

moveDn t@(Tree {selectedIndex = sIx})
  | sIx == treeSize - 1 = t
  | otherwise           = t { selectedIndex = sIx+1 }
  where
    treeSize = length $ concatMap (flatten . filterVisible) $ nodes t


----------------------------------------------------------------------
-- collapse/expand

-- | If tree expanded then collapse it, else expand.
collapseOrExpand t
  = if collapsed $ current t then expand t
                             else collapse t

-- | Collapse or expand only current node.
collapse t@(Tree {..}) = t {nodes = setCollapsed True selectedIndex nodes}
expand t@(Tree {..}) = t {nodes = setCollapsed False selectedIndex nodes}

-- FIXME: работает совсем неправильно
-- | Collapse current node or if already collapsed then collapse
-- parent.
collapse' t@(Tree {..}) = t
  { selectedIndex = parentIx
  , nodes = setCollapsed True parentIx nodes
  }
  where
    -- ищем узел, который нужно свернуть. Это может быть:
    --   - выбранный узел, если у него есть вложенные узлы и он сам
    --     развёрнут
    --   - иначе, сворачиваем родителя выбранного узла
    parentIx = fst $ foldl getParent1 (0, 0) nodes

    -- для самого верхнего уровня логика немного отличается:
    -- если узел свёрнут или он лист, то ничего не происходит (потому
    -- как родительского узла нет)
    getParent1 (p,ix) n = if ix <= selectedIndex then getParent' (ix,ix) n else (p,ix)
    -- у вложенных узлов всегда есть родитель которого можно свернуть
    getParent2 (p,ix) n = if ix <= selectedIndex then getParent' (p,ix) n else (p,ix)

    getParent' (p,ix) (Node n ns)
      -- если выбран лист дерева или свёрнутый узел, нужно сворачивать
      -- ветку на которой он висит
      -- если выбран развёрнутый узел дерева, значит его самого и
      -- сворачиваем
      | ix == selectedIndex = if null ns || collapsed n then (p,ix+1) else (ix,ix+1)
      -- до выбранного узла ещё не добрались, встретили развёрнутый,
      -- запоминаем его как родительский и заходим внутрь
      | (not.collapsed) n && (not.null) ns = foldl getParent2 (ix,ix+1) ns
      -- встретили лист или свёрнутый узел, проходим мимо
      | otherwise = (p,ix+1)

setCollapsed value nodeIx nodes = evalState (mapM setCollapsed' nodes) 0
  where
    setCollapsed' (Node n ns) = do
      ix <- get
      put $ ix + 1
      if ix == nodeIx
          -- проверяем null ns, чтобы не было крестиков около листьев
          then return $ Node (n {collapsed = value && (not.null) ns}) ns
        else if collapsed n || ix > nodeIx
          then return $ Node n ns
        else Node n `fmap` mapM setCollapsed' ns

{- TODO: сделать monad transformer
appendIndex nodes = evalState (mapM (mapNode index) nodes) 0
  where
    index = getIndex >>= \ix -> put (Index $ ix+1) >> return ix
    noIndex = return $ Index (-1)

    mapNode indexProvider (Node n ns) = do
      ix <- indexProvider
      ns' <- mapM (mapNode $ if collapsed n then noIndex else index) ns
      return $ Node (ix,n) ns'
-}


----------------------------------------------------------------------
-- rendering

instance Widget Tree where
    growHorizontal t = isNothing (treeWidth t)
    growVertical   t = isNothing (treeHeight t)
    primaryAttribute _ = def_attr
    withAttribute w _ = w

    render rgn t@(Tree {..})
      = vert_cat
        $  map (renderNode width nodeAttr) a
        ++ [renderNode width selectedAttr n]
        ++ map (renderNode width nodeAttr) b
      where
        width  = maybe (fromIntegral $ region_width rgn)  id treeWidth
        height = maybe (fromIntegral $ region_height rgn) id treeHeight
        Tree {windowStart = start} = resize width height t
        (a, n:b) = splitAt (selectedIndex - start) $ window start height t

renderNode width attr (depth, n) = string attr $ take' width ' ' $ fst $ value n

resize _ h t@(Tree {windowStart = ws, selectedIndex = sIx})
  -- current node is out of view => scroll up until it is visible
  | sIx > ws + h - 1 = t { windowStart = sIx - h + 1 }
  -- we have free space after last node => scroll down to fill it
  | wl < h           = t { windowStart = max 0 $ ws - (h - wl) }
  | otherwise        = t
  where
    wl = length $ window ws h t


---
cur = snd . value . current
current t = (concatMap flatten $ map filterVisible $ nodes t) !! selectedIndex t
window start sz = take sz . drop start . concatMap (flatten . appendDepths 0 . filterVisible) . nodes
appendDepths d (Node n ns) = Node (d,n) $ map (appendDepths (d+1)) ns
filterVisible  (Node n ns) = Node n $ if collapsed n then [] else map filterVisible ns
take' n add lst = if length lst < n
                  then lst ++ replicate (n - length lst) add
                  else take n lst