/*
 * Copyright (c) 2016 Galois, Inc.
 * Distributed under the terms of the BSD3 license (see LICENSE file)
 *
 * This module contains definitions that we wish to eventually promote
 * into the Prelude, but which currently cause typechecking of the
 * Prelude to take too long (see #299)
 */

module Cryptol::Extras where

infixr 5 ==>

/**
 * Logical implication
 */
(==>) : Bit -> Bit -> Bit
a ==> b = if a then b else True

/**
 * Logical negation
 */
not : {a} a -> a
not a = ~ a

/**
 * Conjunction
 */
and : {n} (fin n) => [n]Bit -> Bit
and xs = ~zero == xs

/**
 * Disjunction
 */
or : {n} (fin n) => [n]Bit -> Bit
or xs = zero != xs

/**
 * Conjunction after applying a predicate to all elements.
 */
all : {a,n} (fin n) => (a -> Bit) -> [n]a -> Bit
all f xs = and (map f xs)

/**
 * Disjunction after applying a predicate to all elements.
 */
any : {a,n} (fin n) => (a -> Bit) -> [n]a -> Bit
any f xs = or (map f xs)

/**
 * Map a function over an array.
 */
map : {a, b, n} (a -> b) -> [n]a -> [n]b
map f xs = [f x | x <- xs]

/**
 * Functional left fold.
 *
 * foldl (+) 0 [1,2,3] = ((0 + 1) + 2) + 3
 */
foldl : {a, b, n} (fin n) => (a -> b -> a) -> a -> [n]b -> a
foldl f acc xs = ys ! 0
 where ys = [acc] # [f a x | a <- ys | x <- xs]

/**
 * Functional right fold.
 *
 * foldr (-) 0 [1,2,3] = 0 - (1 - (2 - 3))
 */
foldr : {a,b,n} (fin n) => (a -> b -> b) -> b -> [n]a -> b
foldr f acc xs = ys ! 0
 where ys = [acc] # [f x a | a <- ys | x <- reverse xs]

/**
 * Compute the sum of the words in the array.
 */
sum : {a,n} (fin n, Arith a) => [n]a -> a
sum xs = foldl (+) zero xs

/**
 * Scan left is like a fold that emits the intermediate values.
 */
scanl : {b, a, n}  (b -> a -> b) -> b -> [n]a -> [n+1]b
scanl f acc xs = ys
 where
  ys = [acc] # [f a x | a <- ys | x <- xs]

/**
 * Scan right
 */
scanr : {a,b,n} (fin n) => (a -> b -> b) -> b -> [n]a -> [n+1]b
scanr f acc xs = reverse ys
    where
     ys = [acc] # [f x a | a <- ys | x <- reverse xs]

/**
 * Zero extension
 */
extend : {total,n} (fin total, fin n, total >= n) => [n]Bit -> [total]Bit
extend n = zero # n

/**
 * Signed extension. `extendSigned 0bwxyz : [8] == 0bwwwwwxyz`.
 */
extendSigned : {total,n} (fin total, fin n, n >= 1, total >= n+1) => [n]Bit -> [total]Bit
extendSigned  xs = repeat (xs @ 0) # xs

/**
 * Repeat a value.
 */
repeat : {n, a} a -> [n]a
repeat x = [ x | _ <- zero ]

/**
 * `elem x xs` Returns true if x is equal to a value in xs.
 */
elem : {n,a} (fin n, Cmp a) => a -> [n]a -> Bit
elem a xs = any (\x -> x == a) xs

/**
 * Create a list of tuples from two lists.
 */
zip : {a,b,n} [n]a -> [n]b -> [n](a,b)
zip xs ys = [(x,y) | x <- xs | y <- ys]

/**
 * Create a list by applying the function to each pair of elements in the input.
 * lists
 */
zipWith : {a,b,c,n} (a -> b -> c) -> [n]a -> [n]b -> [n]c
zipWith f xs ys = [f x y | x <- xs | y <- ys]

/**
 * Transform a function into uncurried form.
 */
uncurry : {a,b,c} (a -> b -> c) -> (a,b) -> c
uncurry f = \(a,b) -> f a b

/**
 * Transform a function into curried form.
 */
curry : {a,b,c} ((a, b) -> c) -> a -> b -> c
curry f = \a b -> f (a,b)