{-# LANGUAGE FlexibleContexts #-}


-----------------------------------------------------------------------------
-- |
-- Module      :  Data.Tree.Lens
-- Copyright   :  (C) 2012-16 Edward Kmett
-- License     :  BSD-style (see the file LICENSE)
-- Maintainer  :  Edward Kmett <ekmett@gmail.com>
-- Stability   :  provisional
-- Portability :  MTPCs
--
----------------------------------------------------------------------------

module Data.Tree.Lens
  ( root
  , branches
  ) where

import Prelude ()

import Control.Lens.Internal.Prelude
import Control.Lens
import Data.Tree

-- $setup
-- >>> import Control.Lens
-- >>> import Data.Tree

-- | A 'Lens' that focuses on the root of a 'Tree'.
--
-- >>> view root $ Node 42 []
-- 42
root :: Lens' (Tree a) a
root :: (a -> f a) -> Tree a -> f (Tree a)
root a -> f a
f (Node a
a Forest a
as) = (a -> Forest a -> Tree a
forall a. a -> Forest a -> Tree a
`Node` Forest a
as) (a -> Tree a) -> f a -> f (Tree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f a
a
{-# INLINE root #-}

-- | A 'Lens' returning the direct descendants of the root of a 'Tree'
--
-- @'view' 'branches' ≡ 'subForest'@
branches :: Lens' (Tree a) [Tree a]
branches :: ([Tree a] -> f [Tree a]) -> Tree a -> f (Tree a)
branches [Tree a] -> f [Tree a]
f (Node a
a [Tree a]
as) = a -> [Tree a] -> Tree a
forall a. a -> Forest a -> Tree a
Node a
a ([Tree a] -> Tree a) -> f [Tree a] -> f (Tree a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Tree a] -> f [Tree a]
f [Tree a]
as
{-# INLINE branches #-}