module Unison.Util.List where

import Data.List qualified as List
import Data.List.Extra qualified as List
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 [v]
groupBy :: forall (f :: * -> *) k v.
(Foldable f, Ord k) =>
(v -> k) -> f v -> Map k [v]
groupBy v -> k
f f v
vs = [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] -> v -> Map k [v]) -> Map k [v] -> f 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] -> v -> Map k [v]
step Map k [v]
forall k a. Map k a
Map.empty f v
vs
  where
    step :: Map k [v] -> v -> Map k [v]
step Map k [v]
m v
v = ([v] -> [v] -> [v]) -> k -> [v] -> Map k [v] -> Map k [v]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [v] -> [v] -> [v]
forall a. [a] -> [a] -> [a]
(++) (v -> k
f v
v) [v
v] Map k [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, [b])]
groupMap :: forall (f :: * -> *) k a b.
(Foldable f, Eq k) =>
(a -> (k, b)) -> f a -> [(k, [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, [b])) -> [(k, [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, (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)