module util where

// Determines if a given list is sorted
isSorted : {a, b} (Cmp a, fin b) => [b]a -> Bit
isSorted l = ~0 == [ e1 <= e2
                   | e1 <- ([zero] # l) >> 1
                   | e2 <- ([zero] # l) ]

// Determines if a given list is sorted and all sublists are sorted
isSorted2d : {a, b, c} (Cmp b, fin a, fin c) => [a][c]b -> Bit
isSorted2d l = isSorted l && ~0 == [isSorted e | e <- l]

// Trims or pads a list (usually a word) to fit the type
cast : {a, b, c} (fin a, fin b) => [a]c -> [b]c
cast list = drop ((zero:[b]c) # list)

// Fold from the left
foldl : {a, b, c} (fin c) => (a -> b -> a) -> a -> [c]b -> a
foldl fun base list = intermediates ! 0
  where intermediates = [fun pr el
                        | el <- list
                        | pr <- [base]#intermediates]

// Returns the sum of the elements of a list.
sum : {a, b} (fin a, fin b) => [a][b] -> [b]
sum list = foldl (\x y -> x+y) 0 ([0]#list)

// Converts from ascii single-character decimal digits to
// integers. Note that This has no error checking.
atoi : {a} (fin a) => Char -> [a]
atoi c = cast (c - '0')

// Adds two numbers, adding a bit to ensure no overflow.
safePlus : {a, b} (fin a, fin b) => [a] -> [b] -> [(max a b)+1]
safePlus word1 word2 = (cast word1) + (cast word2)

// Multiplies two numbers, resizing to ensure no overflow.
safeTimes : {a, b} (fin a, fin b) => [a] -> [b] -> [a+b]
safeTimes word1 word2 = (cast word1) * (cast word2)

// Increment a number, adding a bit to ensure no overflow.
inc : {a} (fin a) => [a] -> [a+1]
inc word = (cast word) + 1

// Returns true if x is an element of set, false otherwise.
isElem : {a, b} (Cmp a, fin b) => a -> [b]a -> Bit
isElem x set = 0 != [x == y | y <- set]

// Counts the number of times a given element occurs in list.
countElem : {a, b, c} (Cmp b, fin a, fin c) => b -> [a]b -> [c]
countElem elem list =
  sum [if elem == item then (cast (1:[1])) else 0 | item <- list]

// Counts the number of different elements from set that appear in
// list. This assumes that set is in fact a set (contains no
// duplicates.)
countElemsFromSet : {a, b, c, d} (Cmp c, fin a, fin b, fin d) =>
                    [a]c -> [b]c -> [d]
countElemsFromSet list set =
  sum [if isElem e list then (cast (1:[1])) else 0 | e <- set]

// Returns true if two lists of elements contain the same elements,
// some of which may be duplicates, regardless of order.
isPermutation : {a, b} (Cmp b, fin (width a), a >= 1) =>
                [a]b -> [a]b -> Bool
isPermutation r1 r2 =
  ~0 == [(countElem`{c=width a} e r1) == (countElem e r2) | e <- r1]

// Returns the sum of the elements of a list
safeSum : {a, b} (fin a, fin b) => [a][b] -> [b+a]
safeSum list = foldl (\x y -> (cast x) + (cast y)) 0 ([0]#list)

// Given a list, removes the element indicated by offset, returning
// the remaining list.
removeElem : {a, b, c} (fin b, fin c) => [b] -> [1 + c]a -> [c]a
removeElem offset list = take new_list
  where
    new_list = [if i < offset then normal else shift
               | normal <- list
               | shift <- list << 1
               | i <- [0...]]

// Returns all pairs of elements in a list. Order matters, so each
// pair of elements is included twice in both possible orders.
pairs : {a, b} (fin (1 + a)) => [1 + a]b -> [(1 + a) * a][2]b
pairs list = join ps
  where
    shifts = [list] # [shift <<< 1 | shift <- shifts]
    pairWithEach a alist = [[a, b] | b <- alist]
    ps = [pairWithEach a (tail shift)
         | shift <- shifts
         | a <- list]

// Checks each pair of items from a list with the function comp and
// returns the resulting list of booleans.
pairCompare : {a, b} (fin (1 + a)) =>
              [a + 1]b -> (b -> b -> Bit) -> [(a + 1) * a]
pairCompare list comp =
  [comp elem1 elem2 | [elem1, elem2] <- pairs list]

// returns true if list contains at least n instances of elem
hasNInstances : {a, b, c} (Cmp c, fin b, fin a, a >= 1) =>
                [a] -> c -> [b]c -> Bit
hasNInstances n elem list =
  n == foldl fun 0 list
  where fun c item = if c == n then n
                     else if item == elem then c+1
                          else c

// returns true if list contains at least n different elements of set
hasNElemsFromSet : {a, b, c, d} (Cmp d, fin b, fin c, fin a,
                                 a >= 1) => [a] -> [b]d -> [c]d -> Bit
hasNElemsFromSet n list set =
  n == foldl fun 0 set
  where fun c item = if c == n then n
                     else if isElem item list then c+1
                          else c