--
--
-- Collection
--
--

inductive pattern [a] :=
    | []
    | (::) a [a]
    | (++) [a] [a]
    | (*:) [a] a

--
-- List
--
def list {a} (m: Matcher a) : Matcher [a] :=
  matcher
    | [] as () with
      | [] -> [()]
      | _ -> []
    | $ :: $ as (m, list m) with
      | $x :: $xs -> [(x, xs)]
      | _ -> []
    | $ *: $ as (list m, m) with
      | $xs *: $x -> [(xs, x)]
      | _ -> []
    | _ ++ $ :: _ as (m) with
      | $tgt -> tgt
    | _ ++ $ as (list m) with
      | $tgt ->
        matchAll tgt as list m with
          | loop $i (1, _)
              (_ :: ...)
              $rs -> rs
    | $ ++ $ as (list m, list m) with
      | $tgt ->
        matchAll tgt as list m with
          | loop $i (1, $n)
              ($xa_i :: ...)
              $rs -> (foldr (\i r -> xa_i :: r) [] [1..n], rs)
    | #$val as () with
      | $tgt -> if val = tgt then [()] else []
    | $ as (something) with
      | $tgt -> [tgt]

def sortedList {Ord a} (m: Matcher a) : Matcher [a] :=
  matcher
    | [] as () with
      | [] -> [()]
      | _ -> []
    | $ ++ #$px :: $ as (sortedList m, sortedList m) with
      | $tgt ->
        matchAll tgt as list m with
          | loop $i (1, $n)
              ((?(< px) & $xa_i) :: ...)
              (#px :: $rs) -> (map (\i -> xa_i) [1..n], rs)
    | $ ++ $ as (sortedList m, sortedList m) with
      | $tgt ->
        matchAll tgt as list m with
          | loop $i (1, $n)
              ($xa_i :: ...)
              $rs -> (map (\i -> xa_i) [1..n], rs)
    | $ :: $ as (m, sortedList m) with
      | $x :: $xs -> [(x, xs)]
      | _ -> []
    | #$val as () with
      | $tgt -> if val = tgt then [()] else []
    | $ as (something) with
      | $tgt -> [tgt]

instance {Eq a} Eq [a] where
  (==) xs ys :=
    match (xs, ys) as (list something, list something) with
      | ([], []) -> True
      | ($x :: $xs', $y :: $ys') -> x == y && xs' == ys'
      | _ -> False
  (/=) xs ys := not (xs == ys)

instance {Ord a} Ord [a] where
  compare c1 c2 :=
    match (c1, c2) as (list something, list something) with
      | ([], []) -> Equal
      | ([], _) -> Less
      | (_, []) -> Greater
      | ($x :: $xs, #x :: $ys) -> compare xs ys
      | ($x :: _, $y :: _) -> compare x y
  (<) xs ys := compare xs ys = Less
  (<=) xs ys := compare xs ys != Greater
  (>) xs ys := compare xs ys = Greater
  (>=) xs ys := compare xs ys != Less
  min xs ys := if compare xs ys = Less then xs else ys
  max xs ys := if compare xs ys = Greater then xs else ys

def minimum {Ord a} (xs: [a]) : a :=
  foldl1 min xs

def maximum {Ord a} (xs: [a]) : a :=
  foldl1 max xs

-- splitByOrdering uses $x pattern which doesn't support type annotations
def splitByOrdering {Ord a} (p: a) (xs: [a]) : ([a], [a], [a]) :=
  match xs as list something with
    | [] -> ([], [], [])
    | $x :: $rs ->
        let (ys1, ys2, ys3) := splitByOrdering p rs
         in match compare x p as ordering with
              | less    -> (x :: ys1, ys2, ys3)
              | equal   -> (ys1, x :: ys2, ys3)
              | greater -> (ys1, ys2, x :: ys3)

def sort {Ord a} (xs: [a]) : [a] :=
  match xs as list something with
    | [] -> []
    | $x :: [] -> [x]
    | _ ->
        let n := length xs
            p := nth (i.quotient n 2) xs
            (ys1, ys2, ys3) := splitByOrdering p xs
         in sort ys1 ++ ys2 ++ sort ys3

def merge {Ord a} (xs: [a]) (ys: [a]) : [a] :=
  match (xs, ys) as (list something, list something) with
    | ([], _) -> ys
    | (_, []) -> xs
    | ($x :: $txs, ?(>= x) :: _) -> x :: merge txs ys
    | (_, $y :: $tys) -> y :: merge xs tys

def minimize {a, Ord b} (f: a -> b) (xs: [a]) : a :=
  foldl1 (\x y -> if compare (f x) (f y) = Less then x else y) xs

def maximize {a, Ord b} (f: a -> b) (xs: [a]) : a :=
  foldl1 (\x y -> if compare (f x) (f y) = Greater then x else y) xs


--
-- Accessors
--
def nth {a} (n: Integer) (xs: [a]) : a :=
  match xs as list something with
    | loop $i (1, n - 1, _)
        (_ :: ...)
        ($x :: _) -> x

def takeAndDrop {a} (n: Integer) (xs: [a]) : ([a], [a]) :=
  match xs as list something with
    | loop $i (1, n, _)
        ($a_i :: ...)
        $rs -> (map (\i -> a_i) [1..n], rs)

def take {a} (n: Integer) (xs: [a]) : [a] :=
  if n = 0
    then []
    else match xs as list something with
      | $x :: $xs -> x :: take (n - 1) xs
      | [] -> []

def drop {a} (n: Integer) (xs: [a]) : [a] :=
  if n = 0
    then xs
    else match xs as list something with
      | _ :: $xs -> drop (n - 1) xs
      | [] -> []

def takeWhile {a} (pred: a -> Bool) (xs: [a]) : [a] :=
  match xs as list something with
    | [] -> []
    | $x :: $rs -> if pred x then x :: takeWhile pred rs else []

def takeWhileBy {a} (pred: a -> Bool) (xs: [a]) : [a] :=
  match xs as list something with
    | [] -> []
    | $x :: $rs -> if pred x then x :: takeWhileBy pred rs else [x]

def dropWhile {a} (pred: a -> Bool) (xs: [a]) : [a] :=
  match xs as list something with
    | [] -> []
    | $x :: $rs -> if pred x then dropWhile pred rs else xs

--
-- head, tail, uncons, unsnoc
--
def head {a} (xs: [a]) : a :=
  match xs as list something with
    | $x :: _ -> x

def tail {a} (xs: [a]) : [a] :=
  match xs as list something with
    | _ :: $ys -> ys

def last {a} (xs: [a]) : a :=
  match xs as list something with
    | _ *: $x -> x

def init {a} (xs: [a]) : [a] :=
  match xs as list something with
    | $ys *: _ -> ys

def uncons {a} (xs: [a]) : (a, [a]) :=
  match xs as list something with
    | $x :: $ys -> (x, ys)

def unsnoc {a} (xs: [a]) : ([a], a) :=
  match xs as list something with
    | $ys *: $x -> (ys, x)


--
-- list functions
--
def isEmpty {a} (xs: [a]) : Bool :=
  match xs as list something with
    | [] -> True
    | _  -> False

def length {a} (xs: [a]) : Integer := foldl (\acc _ -> acc + 1) 0 xs

def map {a, b} (fn: a -> b) (xs: [a]) : [b] :=
  match xs as list something with
    | [] -> []
    | $x :: $rs -> fn x :: map fn rs

def map2 {a, b, c} (fn: a -> b -> c) (xs: [a]) (ys: [b]) : [c] :=
  match (xs, ys) as (list something, list something) with
    | ([], _) -> []
    | (_, []) -> []
    | ($x :: $xs2, $y :: $ys2) -> fn x y :: map2 fn xs2 ys2

def map3 {a, b, c, d} (fn: a -> b -> c -> d) (xs: [a]) (ys: [b]) (zs: [c]) : [d] :=
  match (xs, ys, zs) as (list something, list something, list something) with
    | ([], _, _) -> []
    | (_, [], _) -> []
    | (_, _, []) -> []
    | ($x :: $xs2, $y :: $ys2, $z :: $zs2) -> fn x y z :: map3 fn xs2 ys2 zs2

def map4 {a, b, c, d, e} (fn: a -> b -> c -> d -> e) (xs: [a]) (ys: [b]) (zs: [c]) (ws: [d]) : [e] :=
  match (xs, ys, zs, ws) as
    (list something, list something, list something, list something) with
    | ([], _, _, _) -> []
    | (_, [], _, _) -> []
    | (_, _, [], _) -> []
    | (_, _, _, []) -> []
    | ($x :: $xs2, $y :: $ys2, $z :: $zs2, $w :: $ws2) ->
      fn x y z w :: map4 fn xs2 ys2 zs2 ws2

def filter {a} (pred: a -> Bool) (xs: [a]) : [a] := foldr (\y ys -> if pred y then y :: ys else ys) [] xs

def partition {a} (pred: a -> Bool) (xs: [a]) : ([a], [a]) := (filter pred xs, filter 1#(not (pred $1)) xs)

def zip {a, b} (xs: [a]) (ys: [b]) : [(a, b)] := map2 (\x y -> (x, y)) xs ys

def zip3 {a, b, c} (xs: [a]) (ys: [b]) (zs: [c]) : [(a, b, c)] := map3 (\x y z -> (x, y, z)) xs ys zs

def zip4 {a, b, c, d} (xs: [a]) (ys: [b]) (zs: [c]) (ws: [d]) : [(a, b, c, d)] := map4 (\x y z w -> (x, y, z, w)) xs ys zs ws

def lookup {Eq a, b} (k : a) (xs : [(a, b)]) : b :=
  match xs as list (eq, something) with
    | _ ++ (#k, $v) :: _ -> v

def foldr {a, b} (fn : a -> b -> b) (init : b) (ls : [a]) : b :=
  match ls as list something with
    | [] -> init
    | $x :: $xs -> fn x (foldr fn init xs)

def foldl {a, b} (fn : b -> a -> b) (init : b) (ls : [a]) : b :=
  match ls as list something with
    | [] -> init
    | $x :: $xs ->
      let z := fn init x
       in seq z (foldl fn z xs)

def foldl1 {a, b} (fn : b -> a -> b) (ls : [a]) : b := foldl fn (head ls) (tail ls)

def reduce {a, b} (fn : b -> a -> b) (ls : [a]) : b := foldl fn (head ls) (tail ls)

def scanl {a, b} (fn: b -> a -> b) (init: b) (ls: [a]) : [b] :=
  init :: (match ls as list something with
    | [] -> []
    | $x :: $xs -> scanl fn (fn init x) xs)

def iterate {a} (fn: a -> a) (x: a) : [a] :=
  let nx1 := fn x
      nx2 := fn nx1
      nx3 := fn nx2
      nx4 := fn nx3
      nx5 := fn nx4
   in x :: nx1 :: nx2 :: nx3 :: nx4 :: iterate fn nx5

def concat {a} (xss: [[a]]) : [a] := foldr (\xs rs -> xs ++ rs) [] xss

def reverse {a} (xs: [a]) : [a] :=
  match xs as list something with
    | [] -> []
    | $rs *: $x -> x :: reverse rs

def intersperse {a} (sep: a) (ws: [a]) : [a] :=
  match ws as list something with
    | [] -> []
    | $w :: $rs -> foldl (\s1 s2 -> s1 ++ [sep, s2]) [w] rs

def intercalate {a} (sep: [a]) (ws: [[a]]) : [a] := concat (intersperse sep ws)

def split {Eq a} (sep: [a]) (ls: [a]) : [[a]] :=
  match ls as list something with
    | $xs ++ #sep ++ $rs -> xs :: split sep rs
    | _ -> [ls]

def splitAs {a} (m: Matcher a) (sep: [a]) (ls: [a]) : [[a]] :=
  match ls as list m with
    | $xs ++ #sep ++ $rs -> xs :: splitAs m sep rs
    | _ -> [ls]

def splitAt {a} (n: Integer) (ls: [a]) : ([a], [a]) := (take n ls, drop n ls)

def findCycle {a} (xs: [a]) : ([a], [a]) :=
  head
    (matchAll xs as list something with
      | $ys ++ (_ :: _ & $cs) ++ #cs ++ _ -> (ys, cs))

def repeat {a} (xs: [a]) : [a] := xs ++ repeat xs

def repeat1 {a} (x: a) : [a] := x :: repeat1 x

--
-- Others
--
def all {a} (pred: a -> Bool) (xs: [a]) : Bool :=
  match xs as list something with
    | [] -> True
    | $x :: $rs -> if pred x then all pred rs else False

def any {a} (pred: a -> Bool) (xs: [a]) : Bool :=
  match xs as list something with
    | [] -> False
    | $x :: $rs -> if pred x then True else any pred rs

def from (s: Integer) : [Integer] :=
  [s, s + 1, s + 2, s + 3, s + 4, s + 5, s + 6, s + 7, s + 8, s + 9, s + 10] ++
    from (s + 11)

-- Note. `between` is used in the definition of the list matcher.
def between (s: Integer) (e: Integer) : [Integer] :=
  if s = e then [s] else if s < e then s :: between (s + 1) e else []

--
-- Multiset
--
def multiset {a} (m: Matcher a) : Matcher [a] :=
  matcher
    | [] as () with
      | [] -> [()]
      | _ -> []
    | $ :: _ as (m) with
      | $tgt -> tgt
    | $ :: $ as (m, multiset m) with
      | $tgt ->
        matchAll tgt as list m with
          | $hs ++ $x :: $ts -> (x, hs ++ ts)
    | #$pxs ++ $ as (multiset m) with
      | $tgt ->
        match (pxs, tgt) as (list m, multiset m) with
          | loop $i (1, length pxs, _)
              {($x_i :: @, #x_i :: @), ...}
              ([], $rs) -> [rs]
          | _ -> []
    | $ ++ $ as (multiset m, multiset m) with
      | $tgt ->
        matchAll tgt as list m with
          | loop $i (1, $n)
              ($rs_i ++ $x_i :: ...)
              $ts ->
            (map (\i -> x_i) [1..n], concat (map (\i -> rs_i) [1..n] ++ [ts]))
    | #$val as () with
      | $tgt ->
        match (val, tgt) as (list m, multiset m) with
          | ([], []) -> [()]
          | ($x :: $xs, #x :: #xs) -> [()]
          | (_, _) -> []
    | $ as (something) with
      | $tgt -> [tgt]

--
-- multiset operation
--
def deleteFirst {Eq a} (x : a) (xs : [a]) : [a] :=
  match xs as list something with
    | [] -> []
    | #x :: $rs -> rs
    | $y :: $rs -> y :: deleteFirst x rs

def deleteFirstAs {a} (m: Matcher a) (x: a) (xs: [a]) : [a] :=
  match xs as list m with
    | [] -> []
    | #x :: $rs -> rs
    | $y :: $rs -> y :: deleteFirstAs m x rs

def delete {Eq a} (x: a) (xs: [a]) : [a] :=
  match xs as list something with
    | [] -> []
    | $hs ++ #x :: $ts -> hs ++ delete x ts
    | _ -> xs

def deleteAs {a} (m: Matcher a) (x: a) (xs: [a]) : [a] :=
  match xs as list m with
    | [] -> []
    | $hs ++ #x :: $ts -> hs ++ deleteAs m x ts
    | _ -> xs

def difference {Eq a} (xs: [a]) (ys: [a]) : [a] :=
  match ys as list something with
    | [] -> xs
    | $y :: $rs -> difference (deleteFirst y xs) rs

def differenceAs {a} (m: Matcher a) (xs: [a]) (ys: [a]) : [a] :=
  match ys as list m with
    | [] -> xs
    | $y :: $rs -> differenceAs m (deleteFirstAs m y xs) rs

def include {Eq a} (xs: [a]) (ys: [a]) : Bool :=
  match ys as list something with
    | [] -> True
    | $y :: $rs ->
      if member y xs then include (deleteFirst y xs) rs else False

def includeAs {a} (m: Matcher a) (xs: [a]) (ys: [a]) : Bool :=
  match ys as list m with
    | [] -> True
    | $y :: $rs ->
      if memberAs m y xs then includeAs m (deleteFirst y xs) rs else False

def union {Eq a} (xs: [a]) (ys: [a]) : [a] :=
  xs ++ (matchAll (ys, xs) as (multiset something, multiset something) with
    | ($y :: _, !(#y :: _)) -> y)

def unionAs {a} (m: Matcher a) (xs: [a]) (ys: [a]) : [a] :=
  xs ++ (matchAll (ys, xs) as (multiset m, multiset m) with
    | ($y :: _, !(#y :: _)) -> y)

def intersect {Eq a} (xs: [a]) (ys: [a]) : [a] :=
  matchAll (xs, ys) as (multiset something, multiset something) with
    | ($x :: _, #x :: _) -> x

def intersectAs {a} (m: Matcher a) (xs: [a]) (ys: [a]) : [a] :=
  matchAll (xs, ys) as (multiset m, multiset m) with
    | ($x :: _, #x :: _) -> x

--
-- Simple predicate
--
def member {Eq a} (x: a) (ys: [a]) : Bool :=
  match ys as list something with
    | _ ++ #x :: _ -> True
    | _ -> False

def memberAs {a} (m: Matcher a) (x: a) (ys: [a]) : Bool :=
  match ys as list m with
    | _ ++ #x :: _ -> True
    | _ -> False

--
-- Counting
--
def count {Eq a} (x: a) (xs: [a]) : Integer :=
  foldl (\acc y -> if x = y then acc + 1 else acc) 0 xs

def countAs {a} (m: Matcher a) (x: a) (xs: [a]) : Integer :=
  foldl (\acc y -> if eqAs m x y then acc + 1 else acc) 0 xs

def frequency {Eq a} (xs: [a]) : [(a, Integer)] :=
  map (\u -> (u, count u xs)) (unique xs)

def frequencyAs {a} (m: Matcher a) (xs: [a]) : [(a, Integer)] :=
  map (\u -> (u, countAs m u xs)) (uniqueAs m xs)

--
-- Index
--
def elemIndices {Eq a} (x: a) (xs: [a]) : [Integer] :=
  matchAll xs as list something with
    | $hs ++ #x :: _ -> 1 + length hs

--
-- Set
--
def set {a} (m: Matcher a) : Matcher [a] :=
  matcher
    | [] as () with
      | [] -> [()]
      | _ -> []
    | $ :: $ as (m, set m) with
      | $tgt ->
        matchAll tgt as list m with
          | _ ++ $x :: _ -> (x, tgt)
    | #$pxs ++ $ as (set m) with
      | $tgt ->
        match (pxs, tgt) as (list m, set m) with
          | ( loop $i (1, $n) ($x_i :: ...) []
            , loop $i (1, n)  (#x_i :: ...) _ ) -> [tgt]
          | _ -> []
    | $ ++ $ as (set m, set m) with
      | $tgt ->
        matchAll tgt as list m with
          | loop $i (1, $n)
              ($rs_i ++ $x_i :: ...)
              $ts -> (map (\i -> x_i) [1..n], tgt)
    | #$val as () with
      | $tgt ->
        match (unique val, unique tgt) as (list m, multiset m) with
          | ([], []) -> [()]
          | ($x :: $xs, #x :: #xs) -> [()]
          | (_, _) -> []
    | $ as (something) with
      | $tgt -> [tgt]

--
-- set operation
--
def add {Eq a} (x: a) (xs: [a]) : [a] := if member x xs then xs else xs ++ [x]

def addAs {a} (m: Matcher a) (x: a) (xs: [a]) : [a] := if memberAs m x xs then xs else xs ++ [x]

def fastUnique {Eq a} (xs: [a]) : [a] :=
  matchAll sort xs as list something with
    | _ ++ $x :: !(#x :: _) -> x

def unique {Eq a} (xs: [a]) : [a] :=
  reverse
    (matchAll reverse xs as list something with
      | _ ++ $x :: !(_ ++ #x :: _) -> x)

def uniqueAs {a} (m: Matcher a) (xs: [a]) : [a] := loopFn xs []
  where
    loopFn (xs: [a]) (ys: [a]) : [a] :=
      match (xs, ys) as (list m, multiset m) with
        | ([], _) -> ys
        | ($x :: $rs, #x :: _) -> loopFn rs ys
        | ($x :: $rs, _) -> loopFn rs (ys ++ [x])