collection.egi

--
--
-- Collection
--
--

infixr pattern 5 ++

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

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

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

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

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

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

def takeWhile pred xs :=
  match xs as list something with
    | [] -> []
    | $x :: $rs -> if pred x then x :: takeWhile pred rs else []

def takeWhileBy pred xs :=
  match xs as list something with
    | [] -> []
    | $x :: $rs -> if pred x then x :: takeWhileBy pred rs else [x]

def dropWhile pred xs :=
  match xs as list something with
    | [] -> []
    | $x :: $rs -> if pred x then dropWhile pred rs else xs

--
-- head, tail, uncons, unsnoc
--
def head xs :=
  match xs as list something with
    | $x :: _ -> x

def tail xs :=
  match xs as list something with
    | _ :: $ys -> ys

def last xs :=
  match xs as list something with
    | snoc $x _ -> x

def init xs :=
  match xs as list something with
    | snoc _ $ys -> ys

def uncons xs :=
  match xs as list something with
    | $x :: $ys -> (x, ys)

def unsnoc xs :=
  match xs as list something with
    | snoc $x $ys -> (ys, x)


--
-- list functions
--
def isEmpty xs :=
  match xs as list something with
    | [] -> True
    | _  -> False

def length xs := foldl (\acc _ -> acc + 1) 0 xs

def map fn xs :=
  match xs as list something with
    | [] -> []
    | $x :: $rs -> fn x :: map fn rs

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

def map3 fn xs ys zs :=
  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 fn xs ys zs ws :=
  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 pred xs := foldr (\%y %ys -> if pred y then y :: ys else ys) [] xs

def partition pred xs := (filter pred xs, filter 1#(not (pred %1)) xs)

def zip xs ys := map2 (\x y -> (x, y)) xs ys

def zip3 xs ys zs := map3 (\x y z -> (x, y, z)) xs ys zs

def zip4 xs ys zs ws := map4 (\x y z w -> (x, y, z, w)) xs ys zs ws

def lookup k ls :=
  match ls as list (something, something) with
    | _ ++ (#k, $x) :: _ -> x

def foldr fn %init %ls :=
  match ls as list something with
    | [] -> init
    | $x :: $xs -> fn x (foldr fn init xs)

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

def foldl1 fn %ls := foldl fn (head ls) (tail ls)

def reduce fn %ls := foldl fn (head ls) (tail ls)

def scanl fn %init %ls :=
  init :: (match ls as list something with
    | [] -> []
    | $x :: $xs -> scanl fn (fn init x) xs)

def iterate fn %x :=
  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 repeatedSquaring fn %x n :=
  match n as integer with
    | #1 -> x
    | ?isEven ->
      let y := repeatedSquaring fn x (quotient n 2)
       in fn y y
    | ?isOdd ->
      let y := repeatedSquaring fn x (quotient n 2)
       in fn (fn y y) x

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

def reverse xs :=
  match xs as list something with
    | [] -> []
    | snoc $x $rs -> x :: reverse rs

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

def intercalate sep ws := concat (intersperse sep ws)

def split sep ls :=
  match ls as list something with
    | $xs ++ #sep ++ $rs -> xs :: split sep rs
    | _ -> [ls]

def splitAs a sep ls :=
  match ls as list a with
    | $xs ++ #sep ++ $rs -> xs :: splitAs a sep rs
    | _ -> [ls]

def splitAt n ls := (take n ls, drop n ls)

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

def repeat %xs := xs ++ repeat xs

def repeat1 %x := x :: repeat1 x

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

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

def from s :=
  [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 e :=
  if s = e then [s] else if s < e then s :: between (s + 1) e else []

def L./ xs ys :=
  if length xs < length ys
    then ([], xs)
    else match (ys, xs) as (list mathExpr, list mathExpr) with
      | ($y :: $yrs, $x :: $xrs) ->
        let (zs, rs) := L./
                          (map2
                             (-)
                             (take (length yrs) xrs)
                             (map (* (x / y)) yrs) ++ drop (length yrs) xrs)
                          ys
         in (x / y :: zs, rs)

--
-- Multiset
--
def multiset a :=
  matcher
    | [] as () with
      | [] -> [()]
      | _ -> []
    | $ :: _ as (a) with
      | $tgt -> tgt
    | $ :: $ as (a, multiset a) with
      | $tgt ->
        matchAll tgt as list a with
          | $hs ++ $x :: $ts -> (x, hs ++ ts)
    | #$pxs ++ $ as (multiset a) with
      | $tgt ->
        match (pxs, tgt) as (list a, multiset a) with
          | loop $i (1, length pxs, _)
              {($x_i :: @, #x_i :: @), ...}
              ([], $rs) -> [rs]
          | _ -> []
    | $ ++ $ as (multiset a, multiset a) with
      | $tgt ->
        matchAll tgt as list a 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 a, multiset a) with
          | ([], []) -> [()]
          | ($x :: $xs, #x :: #xs) -> [()]
          | (_, _) -> []
    | $ as (something) with
      | $tgt -> [tgt]

--
-- multiset operation
--
def deleteFirst %x xs :=
  match xs as list something with
    | [] -> []
    | #x :: $rs -> rs
    | $y :: $rs -> y :: deleteFirst x rs

def deleteFirstAs a %x xs :=
  match xs as list a with
    | [] -> []
    | #x :: $rs -> rs
    | $y :: $rs -> y :: deleteFirstAs a x rs

def delete x xs :=
  match xs as list something with
    | [] -> []
    | $hs ++ #x :: $ts -> hs ++ delete x ts
    | _ -> xs

def deleteAs a x xs :=
  match xs as list a with
    | [] -> []
    | $hs ++ #x :: $ts -> hs ++ deleteAs a x ts
    | _ -> xs

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

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

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

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

def union xs ys :=
  xs ++ (matchAll (ys, xs) as (multiset something, multiset something) with
    | ($y :: _, !(#y :: _)) -> y)

def unionAs a xs ys :=
  xs ++ (matchAll (ys, xs) as (multiset a, multiset a) with
    | ($y :: _, !(#y :: _)) -> y)

def intersect xs ys :=
  matchAll (xs, ys) as (multiset something, multiset something) with
    | ($x :: _, #x :: _) -> x

def intersectAs a xs ys :=
  matchAll (xs, ys) as (multiset a, multiset a) with
    | ($x :: _, #x :: _) -> x

--
-- Simple predicate
--
def member x ys :=
  match ys as list something with
    | _ ++ #x :: _ -> True
    | _ -> False

def memberAs a x ys :=
  match ys as list a with
    | _ ++ #x :: _ -> True
    | _ -> False

--
-- Counting
--
def count x xs :=
  foldl (\acc y -> if x = y then acc + 1 else acc) 0 xs

def countAs a x xs :=
  foldl (\acc y -> if eqAs a x y then acc + 1 else acc) 0 xs

def frequency xs :=
  map (\u -> (u, count u xs)) (unique xs)

def frequencyAs a xs :=
  map (\u -> (u, countAs a u xs)) (uniqueAs a xs)

--
-- Index
--
def elemIndices x xs :=
  matchAll xs as list something with
    | $hs ++ #x :: _ -> 1 + length hs

--
-- Set
--
def set a :=
  matcher
    | [] as () with
      | [] -> [()]
      | _ -> []
    | $ :: $ as (a, set a) with
      | $tgt ->
        matchAll tgt as list a with
          | _ ++ $x :: _ -> (x, tgt)
    | #$pxs ++ $ as (set a) with
      | $tgt ->
        match (pxs, tgt) as (list a, set a) with
          | ( loop $i (1, $n) ($x_i :: ...) []
            , loop $i (1, n)  (#x_i :: ...) _ ) -> [tgt]
          | _ -> []
    | $ ++ $ as (set a, set a) with
      | $tgt ->
        matchAll tgt as list a 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 a, multiset a) with
          | ([], []) -> [()]
          | ($x :: $xs, #x :: #xs) -> [()]
          | (_, _) -> []
    | $ as (something) with
      | $tgt -> [tgt]

--
-- set operation
--
def add x xs := if member x xs then xs else xs ++ [x]

def addAs a x xs := if memberAs a x xs then xs else xs ++ [x]

def fastUnique xs :=
  matchAll sort xs as list something with
    | _ ++ $x :: !(#x :: _) -> x

def unique xs :=
  reverse
    (matchAll reverse xs as list something with
      | _ ++ $x :: !(_ ++ #x :: _) -> x)

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