{-
Copyright (C) 2011 CE Matthew Farkas-Dyck

This library is free software; it may be modified and/or redistributed under the terms of the GNU Lesser General Public License, as published by the Free Software Foundation, either version 3 of the License, or, optionally, any later version.

This library is distributed in the hope that it may be useful, but WITH NO WARRANTY, not even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.

A copy of the License ought to have come along with this program. Otherwise, see <http://www.gnu.org/licenses/>.
-}

{-# LANGUAGE UndecidableInstances #-}

module Data.LTree (LTree (..), LForest, unfoldLTree, unfoldLTreeM) where

import Prelude hiding (mapM);

import Control.Applicative;
import Control.Monad hiding (mapM);
import Data.Foldable;
import Data.Traversable;

-- |Multary (Rose) Tree, with data solely in leaves
data LTree v a = Stem (LForest v a) | Leaf a;

type LForest v a = v (LTree v a);

instance (Eq a, Eq (LForest v a)) => Eq (LTree v a) where {
  Leaf x  == Leaf y  = x  == y;
  Stem ss == Stem ts = ss == ts;
  _       == _       = False;
};

instance Functor v => Functor (LTree v) where {
    fmap f (Leaf x)  = Leaf (f x);
    fmap f (Stem ts) = Stem (fmap (fmap f) ts);
};

instance Functor v => Applicative (LTree v) where {
  pure x = Leaf x;
  (<*>) (Leaf f)  (Leaf x)  = Leaf (f x);
  (<*>) (Leaf f)  (Stem ts) = Stem (fmap (fmap f) ts);
  (<*>) (Stem ss) t         = Stem (fmap (<*> t) ss);
};

instance Foldable v => Foldable (LTree v) where {
  foldMap f (Stem ts) = foldMap (foldMap f) ts;
  foldMap f (Leaf x)  = f x;
};

instance Traversable v => Traversable (LTree v) where {
  traverse f (Stem ts) = fmap Stem (traverse (traverse f) ts);
  traverse f (Leaf x)  = fmap Leaf (f x);
};

-- Build tree from seed value
unfoldLTree :: Functor v => (b -> Either a (v b)) -> b -> LTree v a;
unfoldLTree f = either Leaf (Stem . fmap (unfoldLTree f)) . f;

--  Build tree from seed value, monadically
unfoldLTreeM :: (Monad m, Traversable v) => (b -> m (Either a (v b))) -> b -> m (LTree v a);
unfoldLTreeM f y = f y >>= either (return . Leaf) (liftM Stem . mapM (unfoldLTreeM f));