Improve complexity of fromMap.

diff --git a/Data/TreeMap/Strict.hs b/Data/TreeMap/Strict.hs
index 08fc7df213e399ce404e021fdb3cbfd9a9bd093e..4ac11951e4a1da05a1543aa9d4bd1ffb2095c8dc 100644 (file)
--- a/Data/TreeMap/Strict.hs
+++ b/Data/TreeMap/Strict.hs
@@ -11,7 +11,7 @@
 -- by 'Path' prefixes (inside a 'Node').
 module Data.TreeMap.Strict where
 
-import Control.Applicative (Applicative(..))
+import Control.Applicative (Applicative(..), Alternative((<|>)))
 import Control.DeepSeq (NFData(..))
 import Control.Monad (Monad(..))
 import Data.Bool
@@ -31,6 +31,7 @@ import Data.Traversable (Traversable(..))
 import Data.Typeable (Typeable)
 import Prelude (Int, Num(..), seq)
 import Text.Show (Show(..))
+import qualified Control.Applicative as App
 import qualified Data.List as List
 import qualified Data.Map.Strict as Map
 import qualified Data.NonNull as NonNull
@@ -49,6 +50,13 @@ instance Semigroup x => Monoid (Strict.Maybe x) where
 instance NFData x => NFData (Strict.Maybe x) where
        rnf Strict.Nothing  = ()
        rnf (Strict.Just x) = rnf x
+instance Applicative Strict.Maybe where
+       pure = Strict.Just
+       Strict.Just f <*> Strict.Just x = Strict.Just (f x)
+       _ <*> _ = Strict.Nothing
+instance Alternative Strict.Maybe where
+       empty = Strict.Nothing
+       x <|> y = if Strict.isJust x then x else y
 
 -- * Type 'TreeMap'
 newtype TreeMap k x
@@ -71,7 +79,6 @@ instance (Ord k, NFData k, NFData x) => NFData (TreeMap k x) where
        rnf (TreeMap m) = rnf m
 
 -- * Type 'Path'
-
 -- | A 'Path' is a non-empty list of 'Map' keys.
 type Path k = NonNull [k]
 
@@ -170,17 +177,28 @@ insert merge p x (TreeMap m) =
                        insert merge p' x node_descendants)
                 k (node Strict.Nothing (insert merge p' x empty)) m
 
--- | Return a 'TreeMap' associating for each tuple of the given list
--- the 'Path' to the value,
--- merging values of identical 'Path's (in respective order).
+-- | Return a 'TreeMap' from a list of 'Path'/value pairs,
+-- with a combining function called on the leftest and rightest values
+-- when their 'Path's are identical.
 fromList :: Ord k => (x -> x -> x) -> [(Path k, x)] -> TreeMap k x
-fromList merge = List.foldl (\acc (p, x) -> insert merge p x acc) empty
+fromList merge = List.foldl' (\acc (p,x) -> insert merge p x acc) empty
 
--- | Return a 'TreeMap' associating for each key and value of the given 'Map'
--- the 'Path' to the value,
--- merging values of identical 'Path's (in respective order).
-fromMap :: Ord k => (x -> x -> x) -> Map (Path k) x -> TreeMap k x
-fromMap merge = Map.foldlWithKey (\acc p x -> insert merge p x acc) empty
+-- | Return a 'TreeMap' from a 'Map' mapping 'Path' to value.
+fromMap :: Ord k => Map (Path k) x -> TreeMap k x
+fromMap = go . Map.toList
+       where
+       go :: Ord k => [(Path k,x)] -> TreeMap k x
+       go m =
+               TreeMap $ Map.fromAscListWith
+                (\Node{node_value=vn, node_descendants=mn}
+                  Node{node_value=vo, node_descendants=mo} ->
+                       node (vn <|> vo) $ union const mn mo) $
+               (<$> m) $ \(p,x) ->
+                       let (p0,mps) = nuncons p in
+                       case mps of
+                        Nothing -> (p0,node (Strict.Just x) empty)
+                        Just ps -> (p0,node Strict.Nothing $ go [(ps,x)])
+-- fromMap = Map.foldlWithKey (\acc p x -> insert const p x acc) empty
 
 -- * Size
 
@@ -247,6 +265,8 @@ union merge (TreeMap tm0) (TreeMap tm1) =
                 (union merge m0 m1))
         tm0 tm1
 
+
+
 -- | Return the 'union' of the given 'TreeMap's.
 --
 -- NOTE: use |List.foldl'| to reduce demand on the control-stack.