module Unison.Util.List where

import Data.List qualified as List
import Data.List.Extra qualified as List
import Data.List.NonEmpty (NonEmpty)
import Data.List.NonEmpty qualified as NEL
import Data.Map qualified as Map
import Data.Set qualified as Set
import Unison.Prelude

multimap :: (Foldable f) => (Ord k) => f (k, v) -> Map k [v]
multimap :: forall (f :: * -> *) k v.
(Foldable f, Ord k) =>
f (k, v) -> Map k [v]
multimap f (k, v)
kvs =
  -- preserve the order of the values from the original list
  [v] -> [v]
forall a. [a] -> [a]
reverse ([v] -> [v]) -> Map k [v] -> Map k [v]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Map k [v] -> (k, v) -> Map k [v])
-> Map k [v] -> f (k, v) -> Map k [v]
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map k [v] -> (k, v) -> Map k [v]
forall {k} {a}. Ord k => Map k [a] -> (k, a) -> Map k [a]
step Map k [v]
forall k a. Map k a
Map.empty f (k, v)
kvs
  where
    step :: Map k [a] -> (k, a) -> Map k [a]
step Map k [a]
m (k
k, a
v) = ([a] -> [a] -> [a]) -> k -> [a] -> Map k [a] -> Map k [a]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++) k
k [a
v] Map k [a]
m

groupBy :: (Foldable f, Ord k) => (v -> k) -> f v -> Map k (NonEmpty v)
groupBy :: forall (f :: * -> *) k v.
(Foldable f, Ord k) =>
(v -> k) -> f v -> Map k (NonEmpty v)
groupBy v -> k
f f v
vs = NonEmpty v -> NonEmpty v
forall a. NonEmpty a -> NonEmpty a
NEL.reverse (NonEmpty v -> NonEmpty v)
-> Map k (NonEmpty v) -> Map k (NonEmpty v)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Map k (NonEmpty v) -> v -> Map k (NonEmpty v))
-> Map k (NonEmpty v) -> f v -> Map k (NonEmpty v)
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map k (NonEmpty v) -> v -> Map k (NonEmpty v)
step Map k (NonEmpty v)
forall k a. Map k a
Map.empty f v
vs
  where
    step :: Map k (NonEmpty v) -> v -> Map k (NonEmpty v)
step Map k (NonEmpty v)
m v
v = (NonEmpty v -> NonEmpty v -> NonEmpty v)
-> k -> NonEmpty v -> Map k (NonEmpty v) -> Map k (NonEmpty v)
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith NonEmpty v -> NonEmpty v -> NonEmpty v
forall a. Semigroup a => a -> a -> a
(<>) (v -> k
f v
v) (v -> NonEmpty v
forall a. a -> NonEmpty a
NEL.singleton v
v) Map k (NonEmpty v)
m

-- | group _consecutive_ elements by a key.
-- e.g.
-- >>> groupMap (\n -> (odd n, show n)) [1, 3, 4, 6, 7]
-- [(True,["1","3"]),(False,["4","6"]),(True,["7"])]
groupMap :: (Foldable f, Eq k) => (a -> (k, b)) -> f a -> [(k, NonEmpty b)]
groupMap :: forall (f :: * -> *) k a b.
(Foldable f, Eq k) =>
(a -> (k, b)) -> f a -> [(k, NonEmpty b)]
groupMap a -> (k, b)
f f a
xs =
  f a
xs
    f a -> (f a -> [a]) -> [a]
forall a b. a -> (a -> b) -> b
& f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList
    [a] -> ([a] -> [(k, b)]) -> [(k, b)]
forall a b. a -> (a -> b) -> b
& (a -> (k, b)) -> [a] -> [(k, b)]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> (k, b)
f
    [(k, b)] -> ([(k, b)] -> [[(k, b)]]) -> [[(k, b)]]
forall a b. a -> (a -> b) -> b
& ((k, b) -> k) -> [(k, b)] -> [[(k, b)]]
forall k a. Eq k => (a -> k) -> [a] -> [[a]]
List.groupOn (k, b) -> k
forall a b. (a, b) -> a
fst
    -- head is okay since groupOn only returns populated lists.
    [[(k, b)]] -> ([(k, b)] -> (k, NonEmpty b)) -> [(k, NonEmpty b)]
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> \[(k, b)]
grp -> ((k, b) -> k
forall a b. (a, b) -> a
fst ((k, b) -> k) -> ([(k, b)] -> (k, b)) -> [(k, b)] -> k
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(k, b)] -> (k, b)
forall a. HasCallStack => [a] -> a
head ([(k, b)] -> k) -> [(k, b)] -> k
forall a b. (a -> b) -> a -> b
$ [(k, b)]
grp, [b] -> NonEmpty b
forall a. HasCallStack => [a] -> NonEmpty a
NEL.fromList ((k, b) -> b
forall a b. (a, b) -> b
snd ((k, b) -> b) -> [(k, b)] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [(k, b)]
grp))

-- returns the subset of `f a` which maps to unique `b`s.
-- prefers earlier copies, if many `a` map to some `b`.
uniqueBy, nubOrdOn :: (Foldable f, Ord b) => (a -> b) -> f a -> [a]
uniqueBy :: forall (f :: * -> *) b a.
(Foldable f, Ord b) =>
(a -> b) -> f a -> [a]
uniqueBy a -> b
f f a
as = [a] -> Set b -> [a]
wrangle' (f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f a
as) Set b
forall a. Set a
Set.empty
  where
    wrangle' :: [a] -> Set b -> [a]
wrangle' [] Set b
_ = []
    wrangle' (a
a : [a]
as) Set b
seen =
      if b -> Set b -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member b
b Set b
seen
        then [a] -> Set b -> [a]
wrangle' [a]
as Set b
seen
        else a
a a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> Set b -> [a]
wrangle' [a]
as (b -> Set b -> Set b
forall a. Ord a => a -> Set a -> Set a
Set.insert b
b Set b
seen)
      where
        b :: b
b = a -> b
f a
a
nubOrdOn :: forall (f :: * -> *) b a.
(Foldable f, Ord b) =>
(a -> b) -> f a -> [a]
nubOrdOn = (a -> b) -> f a -> [a]
forall (f :: * -> *) b a.
(Foldable f, Ord b) =>
(a -> b) -> f a -> [a]
uniqueBy

-- prefers later copies
uniqueBy' :: (Foldable f, Ord b) => (a -> b) -> f a -> [a]
uniqueBy' :: forall (f :: * -> *) b a.
(Foldable f, Ord b) =>
(a -> b) -> f a -> [a]
uniqueBy' a -> b
f = [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> (f a -> [a]) -> f a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> b) -> [a] -> [a]
forall (f :: * -> *) b a.
(Foldable f, Ord b) =>
(a -> b) -> f a -> [a]
uniqueBy a -> b
f ([a] -> [a]) -> (f a -> [a]) -> f a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> (f a -> [a]) -> f a -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList

safeHead :: (Foldable f) => f a -> Maybe a
safeHead :: forall (f :: * -> *) a. Foldable f => f a -> Maybe a
safeHead = [a] -> Maybe a
forall a. [a] -> Maybe a
headMay ([a] -> Maybe a) -> (f a -> [a]) -> f a -> Maybe a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList

validate :: (Semigroup e, Foldable f) => (a -> Either e b) -> f a -> Either e [b]
validate :: forall e (f :: * -> *) a b.
(Semigroup e, Foldable f) =>
(a -> Either e b) -> f a -> Either e [b]
validate a -> Either e b
f f a
as = case [Either e b] -> ([e], [b])
forall a b. [Either a b] -> ([a], [b])
partitionEithers (a -> Either e b
f (a -> Either e b) -> [a] -> [Either e b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> f a -> [a]
forall a. f a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList f a
as) of
  ([], [b]
bs) -> [b] -> Either e [b]
forall a b. b -> Either a b
Right [b]
bs
  (e
e : [e]
es, [b]
_) -> e -> Either e [b]
forall a b. a -> Either a b
Left ((e -> e -> e) -> e -> [e] -> e
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' e -> e -> e
forall a. Semigroup a => a -> a -> a
(<>) e
e [e]
es)

-- Intercalate a list with separators determined by inspecting each
-- adjacent pair.
intercalateMapWith :: (a -> a -> b) -> (a -> b) -> [a] -> [b]
intercalateMapWith :: forall a b. (a -> a -> b) -> (a -> b) -> [a] -> [b]
intercalateMapWith a -> a -> b
sep a -> b
f [a]
xs = [b]
result
  where
    xs' :: [b]
xs' = (a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f [a]
xs
    pairs :: [[a]]
pairs = ([a] -> Bool) -> [[a]] -> [[a]]
forall a. (a -> Bool) -> [a] -> [a]
filter (\[a]
p -> [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
p Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
2) ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$ ([a] -> [a]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
2) ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$ [a] -> [[a]]
forall a. [a] -> [[a]]
List.tails [a]
xs
    seps :: [b]
seps = (([a] -> b) -> [[a]] -> [b]) -> [[a]] -> ([a] -> b) -> [b]
forall a b c. (a -> b -> c) -> b -> a -> c
flip ([a] -> b) -> [[a]] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map [[a]]
pairs (([a] -> b) -> [b]) -> ([a] -> b) -> [b]
forall a b. (a -> b) -> a -> b
$ \case
      a
x1 : a
x2 : [a]
_ -> a -> a -> b
sep a
x1 a
x2
      [a]
_ -> [Char] -> b
forall a. HasCallStack => [Char] -> a
error [Char]
"bad list length"
    paired :: [[b]]
paired = (b -> b -> [b]) -> [b] -> [b] -> [[b]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\b
sep b
x -> [b
sep, b
x]) [b]
seps (Int -> [b] -> [b]
forall a. Int -> [a] -> [a]
drop Int
1 [b]
xs')
    result :: [b]
result = Int -> [b] -> [b]
forall a. Int -> [a] -> [a]
take Int
1 [b]
xs' [b] -> [b] -> [b]
forall a. [a] -> [a] -> [a]
++ [[b]] -> [b]
forall a. Monoid a => [a] -> a
mconcat [[b]]
paired

-- Take runs of consecutive occurrences of r within a list,
-- and in each run, overwrite all but the first occurrence of r with w.
quenchRuns :: (Eq a) => a -> a -> [a] -> [a]
quenchRuns :: forall a. Eq a => a -> a -> [a] -> [a]
quenchRuns a
r a
w = [a] -> [a]
forall a. [a] -> [a]
reverse ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bool -> a -> a -> [a] -> [a] -> [a]
forall {a}. Eq a => Bool -> a -> a -> [a] -> [a] -> [a]
go Bool
False a
r a
w []
  where
    go :: Bool -> a -> a -> [a] -> [a] -> [a]
go Bool
inRun a
r a
w [a]
acc = \case
      [] -> [a]
acc
      a
h : [a]
tl ->
        if a
h a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
r
          then Bool -> a -> a -> [a] -> [a] -> [a]
go Bool
True a
r a
w ((if Bool
inRun then a
w else a
r) a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc) [a]
tl
          else Bool -> a -> a -> [a] -> [a] -> [a]
go Bool
False a
r a
w (a
h a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
acc) [a]
tl

-- | Finds the longest shared path prefix of two paths.
-- Returns (shared prefix, path to first location from shared prefix, path to second location from shared prefix)
--
-- >>> splitOnLongestCommonPrefix ["a", "b", "x"] ["a", "b", "c"]
-- (["a","b"],["x"],["c"])
--
-- >>> splitOnLongestCommonPrefix [] ["a", "b", "c"]
-- ([],[],["a","b","c"])
splitOnLongestCommonPrefix :: (Eq a) => [a] -> [a] -> ([a], [a], [a])
splitOnLongestCommonPrefix :: forall a. Eq a => [a] -> [a] -> ([a], [a], [a])
splitOnLongestCommonPrefix [a]
as [a]
bs =
  case ([a]
as, [a]
bs) of
    ([], [a]
_) -> ([], [a]
as, [a]
bs)
    ([a]
_, []) -> ([], [a]
as, [a]
bs)
    (a
x : [a]
xs, a
y : [a]
ys)
      | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y ->
          let ([a]
prefix, [a]
ra, [a]
rb) = [a] -> [a] -> ([a], [a], [a])
forall a. Eq a => [a] -> [a] -> ([a], [a], [a])
splitOnLongestCommonPrefix [a]
xs [a]
ys
           in (a
x a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
prefix, [a]
ra, [a]
rb)
      | Bool
otherwise -> ([], [a]
as, [a]
bs)