module Unison.Util.Nametree
  ( -- * Nametree
    Nametree (..),
    traverseNametreeWithName,
    unfoldNametree,

    -- ** Flattening and unflattening
    flattenNametree,
    flattenNametrees,
    unflattenNametree,
    unflattenNametrees,
  )
where

import Data.List.NonEmpty (NonEmpty, pattern (:|))
import Data.List.NonEmpty qualified as List.NonEmpty
import Data.Map.Strict qualified as Map
import Data.Semialign (Semialign (alignWith), Unzip (unzipWith), Zip (zipWith))
import Data.These (These (..), these)
import Unison.Name (Name)
import Unison.Name qualified as Name
import Unison.NameSegment (NameSegment)
import Unison.Prelude
import Unison.Util.BiMultimap (BiMultimap)
import Unison.Util.BiMultimap qualified as BiMultimap
import Unison.Util.Defns (Defns (..), DefnsF)
import Prelude hiding (zipWith)

-- | A nametree has a value, and a collection of children nametrees keyed by name segment.
data Nametree a = Nametree
  { forall a. Nametree a -> a
value :: !a,
    forall a. Nametree a -> Map NameSegment (Nametree a)
children :: !(Map NameSegment (Nametree a))
  }
  deriving stock ((forall a b. (a -> b) -> Nametree a -> Nametree b)
-> (forall a b. a -> Nametree b -> Nametree a) -> Functor Nametree
forall a b. a -> Nametree b -> Nametree a
forall a b. (a -> b) -> Nametree a -> Nametree b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
$cfmap :: forall a b. (a -> b) -> Nametree a -> Nametree b
fmap :: forall a b. (a -> b) -> Nametree a -> Nametree b
$c<$ :: forall a b. a -> Nametree b -> Nametree a
<$ :: forall a b. a -> Nametree b -> Nametree a
Functor, (forall m. Monoid m => Nametree m -> m)
-> (forall m a. Monoid m => (a -> m) -> Nametree a -> m)
-> (forall m a. Monoid m => (a -> m) -> Nametree a -> m)
-> (forall a b. (a -> b -> b) -> b -> Nametree a -> b)
-> (forall a b. (a -> b -> b) -> b -> Nametree a -> b)
-> (forall b a. (b -> a -> b) -> b -> Nametree a -> b)
-> (forall b a. (b -> a -> b) -> b -> Nametree a -> b)
-> (forall a. (a -> a -> a) -> Nametree a -> a)
-> (forall a. (a -> a -> a) -> Nametree a -> a)
-> (forall a. Nametree a -> [a])
-> (forall a. Nametree a -> Bool)
-> (forall a. Nametree a -> Int)
-> (forall a. Eq a => a -> Nametree a -> Bool)
-> (forall a. Ord a => Nametree a -> a)
-> (forall a. Ord a => Nametree a -> a)
-> (forall a. Num a => Nametree a -> a)
-> (forall a. Num a => Nametree a -> a)
-> Foldable Nametree
forall a. Eq a => a -> Nametree a -> Bool
forall a. Num a => Nametree a -> a
forall a. Ord a => Nametree a -> a
forall m. Monoid m => Nametree m -> m
forall a. Nametree a -> Bool
forall a. Nametree a -> Int
forall a. Nametree a -> [a]
forall a. (a -> a -> a) -> Nametree a -> a
forall m a. Monoid m => (a -> m) -> Nametree a -> m
forall b a. (b -> a -> b) -> b -> Nametree a -> b
forall a b. (a -> b -> b) -> b -> Nametree a -> b
forall (t :: * -> *).
(forall m. Monoid m => t m -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall m a. Monoid m => (a -> m) -> t a -> m)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall a b. (a -> b -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall b a. (b -> a -> b) -> b -> t a -> b)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. (a -> a -> a) -> t a -> a)
-> (forall a. t a -> [a])
-> (forall a. t a -> Bool)
-> (forall a. t a -> Int)
-> (forall a. Eq a => a -> t a -> Bool)
-> (forall a. Ord a => t a -> a)
-> (forall a. Ord a => t a -> a)
-> (forall a. Num a => t a -> a)
-> (forall a. Num a => t a -> a)
-> Foldable t
$cfold :: forall m. Monoid m => Nametree m -> m
fold :: forall m. Monoid m => Nametree m -> m
$cfoldMap :: forall m a. Monoid m => (a -> m) -> Nametree a -> m
foldMap :: forall m a. Monoid m => (a -> m) -> Nametree a -> m
$cfoldMap' :: forall m a. Monoid m => (a -> m) -> Nametree a -> m
foldMap' :: forall m a. Monoid m => (a -> m) -> Nametree a -> m
$cfoldr :: forall a b. (a -> b -> b) -> b -> Nametree a -> b
foldr :: forall a b. (a -> b -> b) -> b -> Nametree a -> b
$cfoldr' :: forall a b. (a -> b -> b) -> b -> Nametree a -> b
foldr' :: forall a b. (a -> b -> b) -> b -> Nametree a -> b
$cfoldl :: forall b a. (b -> a -> b) -> b -> Nametree a -> b
foldl :: forall b a. (b -> a -> b) -> b -> Nametree a -> b
$cfoldl' :: forall b a. (b -> a -> b) -> b -> Nametree a -> b
foldl' :: forall b a. (b -> a -> b) -> b -> Nametree a -> b
$cfoldr1 :: forall a. (a -> a -> a) -> Nametree a -> a
foldr1 :: forall a. (a -> a -> a) -> Nametree a -> a
$cfoldl1 :: forall a. (a -> a -> a) -> Nametree a -> a
foldl1 :: forall a. (a -> a -> a) -> Nametree a -> a
$ctoList :: forall a. Nametree a -> [a]
toList :: forall a. Nametree a -> [a]
$cnull :: forall a. Nametree a -> Bool
null :: forall a. Nametree a -> Bool
$clength :: forall a. Nametree a -> Int
length :: forall a. Nametree a -> Int
$celem :: forall a. Eq a => a -> Nametree a -> Bool
elem :: forall a. Eq a => a -> Nametree a -> Bool
$cmaximum :: forall a. Ord a => Nametree a -> a
maximum :: forall a. Ord a => Nametree a -> a
$cminimum :: forall a. Ord a => Nametree a -> a
minimum :: forall a. Ord a => Nametree a -> a
$csum :: forall a. Num a => Nametree a -> a
sum :: forall a. Num a => Nametree a -> a
$cproduct :: forall a. Num a => Nametree a -> a
product :: forall a. Num a => Nametree a -> a
Foldable, Functor Nametree
Foldable Nametree
(Functor Nametree, Foldable Nametree) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> Nametree a -> f (Nametree b))
-> (forall (f :: * -> *) a.
    Applicative f =>
    Nametree (f a) -> f (Nametree a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> Nametree a -> m (Nametree b))
-> (forall (m :: * -> *) a.
    Monad m =>
    Nametree (m a) -> m (Nametree a))
-> Traversable Nametree
forall (t :: * -> *).
(Functor t, Foldable t) =>
(forall (f :: * -> *) a b.
 Applicative f =>
 (a -> f b) -> t a -> f (t b))
-> (forall (f :: * -> *) a. Applicative f => t (f a) -> f (t a))
-> (forall (m :: * -> *) a b.
    Monad m =>
    (a -> m b) -> t a -> m (t b))
-> (forall (m :: * -> *) a. Monad m => t (m a) -> m (t a))
-> Traversable t
forall (m :: * -> *) a. Monad m => Nametree (m a) -> m (Nametree a)
forall (f :: * -> *) a.
Applicative f =>
Nametree (f a) -> f (Nametree a)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Nametree a -> m (Nametree b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Nametree a -> f (Nametree b)
$ctraverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Nametree a -> f (Nametree b)
traverse :: forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Nametree a -> f (Nametree b)
$csequenceA :: forall (f :: * -> *) a.
Applicative f =>
Nametree (f a) -> f (Nametree a)
sequenceA :: forall (f :: * -> *) a.
Applicative f =>
Nametree (f a) -> f (Nametree a)
$cmapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Nametree a -> m (Nametree b)
mapM :: forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Nametree a -> m (Nametree b)
$csequence :: forall (m :: * -> *) a. Monad m => Nametree (m a) -> m (Nametree a)
sequence :: forall (m :: * -> *) a. Monad m => Nametree (m a) -> m (Nametree a)
Traversable, (forall x. Nametree a -> Rep (Nametree a) x)
-> (forall x. Rep (Nametree a) x -> Nametree a)
-> Generic (Nametree a)
forall x. Rep (Nametree a) x -> Nametree a
forall x. Nametree a -> Rep (Nametree a) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall a x. Rep (Nametree a) x -> Nametree a
forall a x. Nametree a -> Rep (Nametree a) x
$cfrom :: forall a x. Nametree a -> Rep (Nametree a) x
from :: forall x. Nametree a -> Rep (Nametree a) x
$cto :: forall a x. Rep (Nametree a) x -> Nametree a
to :: forall x. Rep (Nametree a) x -> Nametree a
Generic, Int -> Nametree a -> ShowS
[Nametree a] -> ShowS
Nametree a -> String
(Int -> Nametree a -> ShowS)
-> (Nametree a -> String)
-> ([Nametree a] -> ShowS)
-> Show (Nametree a)
forall a. Show a => Int -> Nametree a -> ShowS
forall a. Show a => [Nametree a] -> ShowS
forall a. Show a => Nametree a -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: forall a. Show a => Int -> Nametree a -> ShowS
showsPrec :: Int -> Nametree a -> ShowS
$cshow :: forall a. Show a => Nametree a -> String
show :: Nametree a -> String
$cshowList :: forall a. Show a => [Nametree a] -> ShowS
showList :: [Nametree a] -> ShowS
Show)

instance Semialign Nametree where
  alignWith :: (These a b -> c) -> Nametree a -> Nametree b -> Nametree c
  alignWith :: forall a b c.
(These a b -> c) -> Nametree a -> Nametree b -> Nametree c
alignWith These a b -> c
f (Nametree a
x Map NameSegment (Nametree a)
xs) (Nametree b
y Map NameSegment (Nametree b)
ys) =
    c -> Map NameSegment (Nametree c) -> Nametree c
forall a. a -> Map NameSegment (Nametree a) -> Nametree a
Nametree (These a b -> c
f (a -> b -> These a b
forall a b. a -> b -> These a b
These a
x b
y)) ((These (Nametree a) (Nametree b) -> Nametree c)
-> Map NameSegment (Nametree a)
-> Map NameSegment (Nametree b)
-> Map NameSegment (Nametree c)
forall a b c.
(These a b -> c)
-> Map NameSegment a -> Map NameSegment b -> Map NameSegment c
forall (f :: * -> *) a b c.
Semialign f =>
(These a b -> c) -> f a -> f b -> f c
alignWith ((Nametree a -> Nametree c)
-> (Nametree b -> Nametree c)
-> (Nametree a -> Nametree b -> Nametree c)
-> These (Nametree a) (Nametree b)
-> Nametree c
forall a c b.
(a -> c) -> (b -> c) -> (a -> b -> c) -> These a b -> c
these ((a -> c) -> Nametree a -> Nametree c
forall a b. (a -> b) -> Nametree a -> Nametree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (These a b -> c
f (These a b -> c) -> (a -> These a b) -> a -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> These a b
forall a b. a -> These a b
This)) ((b -> c) -> Nametree b -> Nametree c
forall a b. (a -> b) -> Nametree a -> Nametree b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (These a b -> c
f (These a b -> c) -> (b -> These a b) -> b -> c
forall b c a. (b -> c) -> (a -> b) -> a -> c
. b -> These a b
forall a b. b -> These a b
That)) ((These a b -> c) -> Nametree a -> Nametree b -> Nametree c
forall a b c.
(These a b -> c) -> Nametree a -> Nametree b -> Nametree c
forall (f :: * -> *) a b c.
Semialign f =>
(These a b -> c) -> f a -> f b -> f c
alignWith These a b -> c
f)) Map NameSegment (Nametree a)
xs Map NameSegment (Nametree b)
ys)

instance Zip Nametree where
  zipWith :: (a -> b -> c) -> Nametree a -> Nametree b -> Nametree c
  zipWith :: forall a b c.
(a -> b -> c) -> Nametree a -> Nametree b -> Nametree c
zipWith a -> b -> c
f (Nametree a
x Map NameSegment (Nametree a)
xs) (Nametree b
y Map NameSegment (Nametree b)
ys) =
    c -> Map NameSegment (Nametree c) -> Nametree c
forall a. a -> Map NameSegment (Nametree a) -> Nametree a
Nametree (a -> b -> c
f a
x b
y) ((Nametree a -> Nametree b -> Nametree c)
-> Map NameSegment (Nametree a)
-> Map NameSegment (Nametree b)
-> Map NameSegment (Nametree c)
forall a b c.
(a -> b -> c)
-> Map NameSegment a -> Map NameSegment b -> Map NameSegment c
forall (f :: * -> *) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith ((a -> b -> c) -> Nametree a -> Nametree b -> Nametree c
forall a b c.
(a -> b -> c) -> Nametree a -> Nametree b -> Nametree c
forall (f :: * -> *) a b c.
Zip f =>
(a -> b -> c) -> f a -> f b -> f c
zipWith a -> b -> c
f) Map NameSegment (Nametree a)
xs Map NameSegment (Nametree b)
ys)

instance Unzip Nametree where
  unzipWith :: (c -> (a, b)) -> Nametree c -> (Nametree a, Nametree b)
  unzipWith :: forall c a b.
(c -> (a, b)) -> Nametree c -> (Nametree a, Nametree b)
unzipWith c -> (a, b)
f (Nametree c
x Map NameSegment (Nametree c)
xs) =
    (a -> Map NameSegment (Nametree a) -> Nametree a
forall a. a -> Map NameSegment (Nametree a) -> Nametree a
Nametree a
y Map NameSegment (Nametree a)
ys, b -> Map NameSegment (Nametree b) -> Nametree b
forall a. a -> Map NameSegment (Nametree a) -> Nametree a
Nametree b
z Map NameSegment (Nametree b)
zs)
    where
      (a
y, b
z) = c -> (a, b)
f c
x
      (Map NameSegment (Nametree a)
ys, Map NameSegment (Nametree b)
zs) = (Nametree c -> (Nametree a, Nametree b))
-> Map NameSegment (Nametree c)
-> (Map NameSegment (Nametree a), Map NameSegment (Nametree b))
forall c a b.
(c -> (a, b))
-> Map NameSegment c -> (Map NameSegment a, Map NameSegment b)
forall (f :: * -> *) c a b.
Unzip f =>
(c -> (a, b)) -> f c -> (f a, f b)
unzipWith ((c -> (a, b)) -> Nametree c -> (Nametree a, Nametree b)
forall c a b.
(c -> (a, b)) -> Nametree c -> (Nametree a, Nametree b)
forall (f :: * -> *) c a b.
Unzip f =>
(c -> (a, b)) -> f c -> (f a, f b)
unzipWith c -> (a, b)
f) Map NameSegment (Nametree c)
xs

-- | Traverse over a nametree, with access to the list of name segments (in reverse order) leading to each value.
traverseNametreeWithName :: (Applicative f) => ([NameSegment] -> a -> f b) -> Nametree a -> f (Nametree b)
traverseNametreeWithName :: forall (f :: * -> *) a b.
Applicative f =>
([NameSegment] -> a -> f b) -> Nametree a -> f (Nametree b)
traverseNametreeWithName [NameSegment] -> a -> f b
f =
  [NameSegment] -> Nametree a -> f (Nametree b)
go []
  where
    go :: [NameSegment] -> Nametree a -> f (Nametree b)
go [NameSegment]
names (Nametree a
x Map NameSegment (Nametree a)
xs) =
      b -> Map NameSegment (Nametree b) -> Nametree b
forall a. a -> Map NameSegment (Nametree a) -> Nametree a
Nametree (b -> Map NameSegment (Nametree b) -> Nametree b)
-> f b -> f (Map NameSegment (Nametree b) -> Nametree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [NameSegment] -> a -> f b
f [NameSegment]
names a
x f (Map NameSegment (Nametree b) -> Nametree b)
-> f (Map NameSegment (Nametree b)) -> f (Nametree b)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (NameSegment -> Nametree a -> f (Nametree b))
-> Map NameSegment (Nametree a) -> f (Map NameSegment (Nametree b))
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
Map.traverseWithKey (\NameSegment
name -> [NameSegment] -> Nametree a -> f (Nametree b)
go (NameSegment
name NameSegment -> [NameSegment] -> [NameSegment]
forall a. a -> [a] -> [a]
: [NameSegment]
names)) Map NameSegment (Nametree a)
xs

-- | Build a nametree from a seed value.
unfoldNametree :: (a -> (b, Map NameSegment a)) -> a -> Nametree b
unfoldNametree :: forall a b. (a -> (b, Map NameSegment a)) -> a -> Nametree b
unfoldNametree a -> (b, Map NameSegment a)
f a
x =
  let (b
y, Map NameSegment a
ys) = a -> (b, Map NameSegment a)
f a
x
   in b -> Map NameSegment (Nametree b) -> Nametree b
forall a. a -> Map NameSegment (Nametree a) -> Nametree a
Nametree b
y ((a -> (b, Map NameSegment a)) -> a -> Nametree b
forall a b. (a -> (b, Map NameSegment a)) -> a -> Nametree b
unfoldNametree a -> (b, Map NameSegment a)
f (a -> Nametree b)
-> Map NameSegment a -> Map NameSegment (Nametree b)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Map NameSegment a
ys)

-- | 'flattenNametree' organizes a nametree like
--
-- > "foo" = #foo
-- > "foo": {
-- >   "bar" = #bar
-- >   "bar": {
-- >     "baz" = #baz
-- >   }
-- > }
--
-- into an equivalent-but-flat association between names and definitions, like
--
-- > {
-- >   "foo" = #bar,
-- >   "foo.bar" = #bar,
-- >   "foo.bar.baz" = #baz
-- > }
flattenNametree ::
  forall a b.
  (Ord b) =>
  (a -> Map NameSegment b) ->
  Nametree a ->
  BiMultimap b Name
flattenNametree :: forall a b.
Ord b =>
(a -> Map NameSegment b) -> Nametree a -> BiMultimap b Name
flattenNametree a -> Map NameSegment b
f =
  [NameSegment] -> Nametree a -> BiMultimap b Name
go []
  where
    go :: [NameSegment] -> Nametree a -> BiMultimap b Name
    go :: [NameSegment] -> Nametree a -> BiMultimap b Name
go [NameSegment]
prefix (Nametree a
node Map NameSegment (Nametree a)
children) =
      ((NameSegment, Nametree a)
 -> BiMultimap b Name -> BiMultimap b Name)
-> BiMultimap b Name
-> [(NameSegment, Nametree a)]
-> BiMultimap b Name
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr
        ( \(NameSegment
name, Nametree a
child) ->
            -- This union is safe because the keys are disjoint
            BiMultimap b Name -> BiMultimap b Name -> BiMultimap b Name
forall a b.
(Ord a, Ord b) =>
BiMultimap a b -> BiMultimap a b -> BiMultimap a b
BiMultimap.unsafeUnion ([NameSegment] -> Nametree a -> BiMultimap b Name
go (NameSegment
name NameSegment -> [NameSegment] -> [NameSegment]
forall a. a -> [a] -> [a]
: [NameSegment]
prefix) Nametree a
child)
        )
        ( Map Name b -> BiMultimap b Name
forall a b. (Ord a, Ord b) => Map b a -> BiMultimap a b
BiMultimap.fromRange
            ( (NameSegment -> Name) -> Map NameSegment b -> Map Name b
forall k1 k2 a. (k1 -> k2) -> Map k1 a -> Map k2 a
Map.mapKeysMonotonic
                (\NameSegment
name -> NonEmpty NameSegment -> Name
Name.fromReverseSegments (NameSegment
name NameSegment -> [NameSegment] -> NonEmpty NameSegment
forall a. a -> [a] -> NonEmpty a
:| [NameSegment]
prefix))
                (a -> Map NameSegment b
f a
node)
            )
        )
        (Map NameSegment (Nametree a) -> [(NameSegment, Nametree a)]
forall k a. Map k a -> [(k, a)]
Map.toList Map NameSegment (Nametree a)
children)

-- | Like 'flattenNametree', but works on both the types and terms namespace at once.
flattenNametrees ::
  (Ord term, Ord typ) =>
  Nametree (DefnsF (Map NameSegment) term typ) ->
  Defns (BiMultimap term Name) (BiMultimap typ Name)
flattenNametrees :: forall term typ.
(Ord term, Ord typ) =>
Nametree (DefnsF (Map NameSegment) term typ)
-> Defns (BiMultimap term Name) (BiMultimap typ Name)
flattenNametrees Nametree (Defns (Map NameSegment term) (Map NameSegment typ))
defns =
  Defns
    { $sel:terms:Defns :: BiMultimap term Name
terms = (Defns (Map NameSegment term) (Map NameSegment typ)
 -> Map NameSegment term)
-> Nametree (Defns (Map NameSegment term) (Map NameSegment typ))
-> BiMultimap term Name
forall a b.
Ord b =>
(a -> Map NameSegment b) -> Nametree a -> BiMultimap b Name
flattenNametree (Getting
  (Map NameSegment term)
  (Defns (Map NameSegment term) (Map NameSegment typ))
  (Map NameSegment term)
-> Defns (Map NameSegment term) (Map NameSegment typ)
-> Map NameSegment term
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting
  (Map NameSegment term)
  (Defns (Map NameSegment term) (Map NameSegment typ))
  (Map NameSegment term)
#terms) Nametree (Defns (Map NameSegment term) (Map NameSegment typ))
defns,
      $sel:types:Defns :: BiMultimap typ Name
types = (Defns (Map NameSegment term) (Map NameSegment typ)
 -> Map NameSegment typ)
-> Nametree (Defns (Map NameSegment term) (Map NameSegment typ))
-> BiMultimap typ Name
forall a b.
Ord b =>
(a -> Map NameSegment b) -> Nametree a -> BiMultimap b Name
flattenNametree (Getting
  (Map NameSegment typ)
  (Defns (Map NameSegment term) (Map NameSegment typ))
  (Map NameSegment typ)
-> Defns (Map NameSegment term) (Map NameSegment typ)
-> Map NameSegment typ
forall s (m :: * -> *) a. MonadReader s m => Getting a s a -> m a
view Getting
  (Map NameSegment typ)
  (Defns (Map NameSegment term) (Map NameSegment typ))
  (Map NameSegment typ)
#types) Nametree (Defns (Map NameSegment term) (Map NameSegment typ))
defns
    }

-- | 'unflattenNametree' organizes an association between names and definitions like
--
-- > {
-- >   "foo" = #bar,
-- >   "foo.bar" = #bar,
-- >   "foo.bar.baz" = #baz
-- > }
--
-- into an equivalent-but-less-flat nametree, like
--
-- > "foo" = #foo
-- > "foo": {
-- >   "bar" = #bar
-- >   "bar": {
-- >     "baz" = #baz
-- >   }
-- > }
unflattenNametree :: (Ord a) => Map Name a -> Nametree (Map NameSegment a)
unflattenNametree :: forall a. Ord a => Map Name a -> Nametree (Map NameSegment a)
unflattenNametree =
  ([(NonEmpty NameSegment, a)]
 -> (Map NameSegment a,
     Map NameSegment [(NonEmpty NameSegment, a)]))
-> [(NonEmpty NameSegment, a)] -> Nametree (Map NameSegment a)
forall a b. (a -> (b, Map NameSegment a)) -> a -> Nametree b
unfoldNametree [(NonEmpty NameSegment, a)]
-> (Map NameSegment a, Map NameSegment [(NonEmpty NameSegment, a)])
forall a.
[(NonEmpty NameSegment, a)]
-> (Map NameSegment a, Map NameSegment [(NonEmpty NameSegment, a)])
unflattenLevel ([(NonEmpty NameSegment, a)] -> Nametree (Map NameSegment a))
-> (Map Name a -> [(NonEmpty NameSegment, a)])
-> Map Name a
-> Nametree (Map NameSegment a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Name, a) -> (NonEmpty NameSegment, a))
-> [(Name, a)] -> [(NonEmpty NameSegment, a)]
forall a b. (a -> b) -> [a] -> [b]
map ((Name -> NonEmpty NameSegment)
-> (Name, a) -> (NonEmpty NameSegment, a)
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first Name -> NonEmpty NameSegment
Name.segments) ([(Name, a)] -> [(NonEmpty NameSegment, a)])
-> (Map Name a -> [(Name, a)])
-> Map Name a
-> [(NonEmpty NameSegment, a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Name a -> [(Name, a)]
forall k a. Map k a -> [(k, a)]
Map.toList
  where
    unflattenLevel :: [(NonEmpty NameSegment, a)] -> (Map NameSegment a, Map NameSegment [(NonEmpty NameSegment, a)])
    unflattenLevel :: forall a.
[(NonEmpty NameSegment, a)]
-> (Map NameSegment a, Map NameSegment [(NonEmpty NameSegment, a)])
unflattenLevel =
      ((Map NameSegment a, Map NameSegment [(NonEmpty NameSegment, a)])
 -> (NonEmpty NameSegment, a)
 -> (Map NameSegment a,
     Map NameSegment [(NonEmpty NameSegment, a)]))
-> (Map NameSegment a, Map NameSegment [(NonEmpty NameSegment, a)])
-> [(NonEmpty NameSegment, a)]
-> (Map NameSegment a, Map NameSegment [(NonEmpty NameSegment, a)])
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Map NameSegment a, Map NameSegment [(NonEmpty NameSegment, a)])
-> (NonEmpty NameSegment, a)
-> (Map NameSegment a, Map NameSegment [(NonEmpty NameSegment, a)])
forall {k} {b}.
Ord k =>
(Map k b, Map k [(NonEmpty k, b)])
-> (NonEmpty k, b) -> (Map k b, Map k [(NonEmpty k, b)])
phi (Map NameSegment a
forall k a. Map k a
Map.empty, Map NameSegment [(NonEmpty NameSegment, a)]
forall k a. Map k a
Map.empty)
      where
        phi :: (Map k b, Map k [(NonEmpty k, b)])
-> (NonEmpty k, b) -> (Map k b, Map k [(NonEmpty k, b)])
phi (!Map k b
accValue, !Map k [(NonEmpty k, b)]
accChildren) = \case
          (NameHere k
n, b
v) -> (k -> b -> Map k b -> Map k b
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert k
n b
v Map k b
accValue, Map k [(NonEmpty k, b)]
accChildren)
          (NameThere k
n NonEmpty k
ns, b
v) -> (Map k b
accValue, ([(NonEmpty k, b)] -> [(NonEmpty k, b)] -> [(NonEmpty k, b)])
-> k
-> [(NonEmpty k, b)]
-> Map k [(NonEmpty k, b)]
-> Map k [(NonEmpty k, b)]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [(NonEmpty k, b)] -> [(NonEmpty k, b)] -> [(NonEmpty k, b)]
forall a. [a] -> [a] -> [a]
(++) k
n [(NonEmpty k
ns, b
v)] Map k [(NonEmpty k, b)]
accChildren)

-- | Like 'unflattenNametree', but works on both the types and terms namespace at once.
unflattenNametrees :: (Ord term, Ord typ) => DefnsF (Map Name) term typ -> Nametree (DefnsF (Map NameSegment) term typ)
unflattenNametrees :: forall term typ.
(Ord term, Ord typ) =>
DefnsF (Map Name) term typ
-> Nametree (DefnsF (Map NameSegment) term typ)
unflattenNametrees DefnsF (Map Name) term typ
defns =
  (These (Map NameSegment term) (Map NameSegment typ)
 -> DefnsF (Map NameSegment) term typ)
-> Nametree (Map NameSegment term)
-> Nametree (Map NameSegment typ)
-> Nametree (DefnsF (Map NameSegment) term typ)
forall a b c.
(These a b -> c) -> Nametree a -> Nametree b -> Nametree c
forall (f :: * -> *) a b c.
Semialign f =>
(These a b -> c) -> f a -> f b -> f c
alignWith
    ( \case
        This Map NameSegment term
terms -> Defns {Map NameSegment term
$sel:terms:Defns :: Map NameSegment term
terms :: Map NameSegment term
terms, $sel:types:Defns :: Map NameSegment typ
types = Map NameSegment typ
forall k a. Map k a
Map.empty}
        That Map NameSegment typ
types -> Defns {$sel:terms:Defns :: Map NameSegment term
terms = Map NameSegment term
forall k a. Map k a
Map.empty, Map NameSegment typ
$sel:types:Defns :: Map NameSegment typ
types :: Map NameSegment typ
types}
        These Map NameSegment term
terms Map NameSegment typ
types -> Defns {Map NameSegment term
$sel:terms:Defns :: Map NameSegment term
terms :: Map NameSegment term
terms, Map NameSegment typ
$sel:types:Defns :: Map NameSegment typ
types :: Map NameSegment typ
types}
    )
    (Map Name term -> Nametree (Map NameSegment term)
forall a. Ord a => Map Name a -> Nametree (Map NameSegment a)
unflattenNametree DefnsF (Map Name) term typ
defns.terms)
    (Map Name typ -> Nametree (Map NameSegment typ)
forall a. Ord a => Map Name a -> Nametree (Map NameSegment a)
unflattenNametree DefnsF (Map Name) term typ
defns.types)

-- Helper patterns for switching on "name here" (1 name segment) or "name there" (2+ name segments)

pattern NameHere :: a -> NonEmpty a
pattern $mNameHere :: forall {r} {a}. NonEmpty a -> (a -> r) -> ((# #) -> r) -> r
NameHere x <- x :| (List.NonEmpty.nonEmpty -> Nothing)

pattern NameThere :: a -> NonEmpty a -> NonEmpty a
pattern $mNameThere :: forall {r} {a}.
NonEmpty a -> (a -> NonEmpty a -> r) -> ((# #) -> r) -> r
NameThere x xs <- x :| (List.NonEmpty.nonEmpty -> Just xs)

{-# COMPLETE NameHere, NameThere #-}