path: root/stdlib/source/library/lux/data/collection/list.lux

(.module:
  [library
   [lux #*
    ["@" target]
    [abstract
     [monoid (#+ Monoid)]
     [apply (#+ Apply)]
     [equivalence (#+ Equivalence)]
     [hash (#+ Hash)]
     [fold (#+ Fold)]
     [predicate (#+ Predicate)]
     ["." functor (#+ Functor)]
     ["." monad (#+ do Monad)]
     ["." enum]]
    [data
     ["." bit]
     ["." product]]
    [math
     [number
      ["n" nat]]]]])

... (type: (List a)
...   #End
...   (#Item a (List a)))

(implementation: .public fold
  (Fold List)
  
  (def: (fold f init xs)
    (case xs
      #.End
      init

      (#.Item x xs')
      (fold f (f x init) xs'))))

(def: .public (folds f init inputs)
  (All [a b] (-> (-> a b b) b (List a) (List b)))
  (case inputs
    #.End
    (list init)
    
    (#.Item [head tail])
    (#.Item [init (folds f (f head init) tail)])))

(def: .public (reversed xs)
  (All [a]
    (-> (List a) (List a)))
  (fold (function (_ head tail) (#.Item head tail))
        #.End
        xs))

(def: .public (only keep? xs)
  {#.doc (doc "A list with only values that satisfy the predicate.")}
  (All [a]
    (-> (Predicate a) (List a) (List a)))
  (case xs
    #.End
    #.End
    
    (#.Item x xs')
    (if (keep? x)
      (#.Item x (only keep? xs'))
      (only keep? xs'))))

(def: .public (partition satisfies? list)
  {#.doc "Divide the list into all elements that satisfy a predicate, and all elements that do not."}
  (All [a] (-> (Predicate a) (List a) [(List a) (List a)]))
  (case list
    #.End
    [#.End #.End]

    (#.Item head tail)
    (let [[in out] (partition satisfies? tail)]
      (if (satisfies? head)
        [(#.Item head in) out]
        [in (#.Item head out)]))))

(def: .public (pairs xs)
  {#.doc (doc "Cut the list into pairs of 2."
              "Caveat emptor: If the list has an un-even number of elements, the last one will be skipped.")}
  (All [a] (-> (List a) (List [a a])))
  (case xs
    (^ (list& x1 x2 xs'))
    (#.Item [x1 x2] (pairs xs'))

    _
    #.End))

(template [<name> <then> <else>]
  [(def: .public (<name> n xs)
     (All [a]
       (-> Nat (List a) (List a)))
     (if (n.> 0 n)
       (case xs
         #.End
         #.End
         
         (#.Item x xs')
         <then>)
       <else>))]
  
  [take (#.Item x (take (dec n) xs')) #.End]
  [drop (drop (dec n) xs') xs]
  )

(template [<name> <then> <else>]
  [(def: .public (<name> predicate xs)
     (All [a]
       (-> (Predicate a) (List a) (List a)))
     (case xs
       #.End
       #.End
       
       (#.Item x xs')
       (if (predicate x)
         <then>
         <else>)))]

  [take_while (#.Item x (take_while predicate xs')) #.End]
  [drop_while (drop_while predicate xs') xs]
  )

(def: .public (split n xs)
  (All [a]
    (-> Nat (List a) [(List a) (List a)]))
  (if (n.> 0 n)
    (case xs
      #.End
      [#.End #.End]
      
      (#.Item x xs')
      (let [[tail rest] (split (dec n) xs')]
        [(#.Item x tail) rest]))
    [#.End xs]))

(def: (split_with' predicate ys xs)
  (All [a]
    (-> (Predicate a) (List a) (List a) [(List a) (List a)]))
  (case xs
    #.End
    [ys xs]

    (#.Item x xs')
    (if (predicate x)
      (split_with' predicate (#.Item x ys) xs')
      [ys xs])))

(def: .public (split_with predicate xs)
  {#.doc "Segment the list by using a predicate to tell when to cut."}
  (All [a]
    (-> (Predicate a) (List a) [(List a) (List a)]))
  (let [[ys' xs'] (split_with' predicate #.End xs)]
    [(reversed ys') xs']))

(def: .public (chunk size list)
  {#.doc "Segment the list in chunks of the given size."}
  (All [a] (-> Nat (List a) (List (List a))))
  (case list
    #.End
    #.End

    _
    (let [[pre post] (split size list)]
      (#.Item pre (chunk size post)))))

(def: .public (repeated n x)
  {#.doc "A list of the value x, repeated n times."}
  (All [a]
    (-> Nat a (List a)))
  (case n
    0 #.End
    _ (#.Item x (repeated (dec n) x))))

(def: (iterations' f x)
  (All [a]
    (-> (-> a (Maybe a)) a (List a)))
  (case (f x)
    (#.Some x')
    (#.Item x (iterations' f x'))

    #.None
    (list)))

(def: .public (iterations f x)
  {#.doc "Generates a list element by element until the function returns #.None."}
  (All [a]
    (-> (-> a (Maybe a)) a (List a)))
  (case (f x)
    (#.Some x')
    (#.Item x (iterations' f x'))

    #.None
    (list x)))

(def: .public (one check xs)
  (All [a b]
    (-> (-> a (Maybe b)) (List a) (Maybe b)))
  (case xs
    #.End
    #.None

    (#.Item x xs')
    (case (check x)
      (#.Some output)
      (#.Some output)
      
      #.None
      (one check xs'))))

(def: .public (all check xs)
  (All [a b]
    (-> (-> a (Maybe b)) (List a) (List b)))
  (for {... TODO: Stop relying on this ASAP.
        @.js
        (fold (function (_ head tail)
                (case (check head)
                  (#.Some head)
                  (#.Item head tail)
                  
                  #.None
                  tail))
              #.End
              (reversed xs))}
       (case xs
         #.End
         #.End

         (#.Item x xs')
         (case (check x)
           (#.Some output)
           (#.Item output (all check xs'))
           
           #.None
           (all check xs')))))

(def: .public (find predicate xs)
  {#.doc "Yields the first value in the list that satisfies the predicate."}
  (All [a]
    (-> (Predicate a) (List a) (Maybe a)))
  (..one (function (_ value)
           (if (predicate value)
             (#.Some value)
             #.None))
         xs))

(def: .public (interpose sep xs)
  {#.doc "Puts a value between every two elements in the list."}
  (All [a]
    (-> a (List a) (List a)))
  (case xs
    #.End
    xs

    (#.Item x #.End)
    xs

    (#.Item x xs')
    (list& x sep (interpose sep xs'))))

(def: .public (size list)
  (All [a] (-> (List a) Nat))
  (fold (function (_ _ acc) (n.+ 1 acc)) 0 list))

(template [<name> <init> <op>]
  [(def: .public (<name> predicate xs)
     (All [a]
       (-> (Predicate a) (List a) Bit))
     (loop [xs xs]
       (case xs
         #.End
         <init>

         (#.Item x xs')
         (case (predicate x)
           <init>
           (recur xs')

           output
           output))))]

  [every? #1 and]
  [any?   #0 or]
  )

(def: .public (item i xs)
  {#.doc "Fetches the element at the specified index."}
  (All [a]
    (-> Nat (List a) (Maybe a)))
  (case xs
    #.End
    #.None

    (#.Item x xs')
    (if (n.= 0 i)
      (#.Some x)
      (item (dec i) xs'))))

(implementation: .public (equivalence Equivalence<a>)
  (All [a] (-> (Equivalence a) (Equivalence (List a))))
  
  (def: (= xs ys)
    (case [xs ys]
      [#.End #.End]
      #1

      [(#.Item x xs') (#.Item y ys')]
      (and (\ Equivalence<a> = x y)
           (= xs' ys'))

      [_ _]
      #0
      )))

(implementation: .public (hash super)
  (All [a] (-> (Hash a) (Hash (List a))))

  (def: &equivalence
    (..equivalence (\ super &equivalence)))
  
  (def: hash
    (\ ..fold fold
       (function (_ member hash)
         (n.+ (\ super hash member) hash))
       0)))

(implementation: .public monoid
  (All [a] (Monoid (List a)))
  
  (def: identity #.End)
  (def: (compose xs ys)
    (case xs
      #.End
      ys
      
      (#.Item x xs')
      (#.Item x (compose xs' ys)))))

(open: "." ..monoid)

(implementation: .public functor
  (Functor List)
  
  (def: (map f ma)
    (case ma
      #.End
      #.End
      
      (#.Item a ma')
      (#.Item (f a) (map f ma')))))

(open: "." ..functor)

(implementation: .public apply
  (Apply List)
  
  (def: &functor ..functor)

  (def: (apply ff fa)
    (case ff
      #.End
      #.End
      
      (#.Item f ff')
      (compose (map f fa) (apply ff' fa)))))

(implementation: .public monad
  (Monad List)
  
  (def: &functor ..functor)

  (def: (in a)
    (#.Item a #.End))

  (def: join
    (|>> reversed (fold compose identity))))

(def: .public (sort < xs)
  {#.doc (doc "A list ordered by a comparison function.")}
  (All [a] (-> (-> a a Bit) (List a) (List a)))
  (case xs
    #.End
    (list)
    
    (#.Item x xs')
    (let [[pre post] (fold (function (_ x' [pre post])
                             (if (< x x')
                               [(#.Item x' pre) post]
                               [pre (#.Item x' post)]))
                           [(list) (list)]
                           xs')]
      ($_ compose (sort < pre) (list x) (sort < post)))))

(def: .public (empty? xs)
  (All [a] (Predicate (List a)))
  (case xs
    #.End
    true
    
    _
    false))

(def: .public (member? eq xs x)
  (All [a] (-> (Equivalence a) (List a) a Bit))
  (case xs
    #.End
    #0
    
    (#.Item x' xs')
    (or (\ eq = x x')
        (member? eq xs' x))))

(template [<name> <output> <side> <doc>]
  [(def: .public (<name> xs)
     {#.doc <doc>}
     (All [a] (-> (List a) (Maybe <output>)))
     (case xs
       #.End
       #.None

       (#.Item x xs')
       (#.Some <side>)))]

  [head a        x   "Yields the first element of a list."]
  [tail (List a) xs' "For a list of size N, yields the N-1 elements after the first one."]
  )

(def: .public (indices size)
  {#.doc "Produces all the valid indices for a given size."}
  (All [a] (-> Nat (List Nat)))
  (if (n.= 0 size)
    (list)
    (|> size dec (enum.range n.enum 0))))

(def: (identifier$ name)
  (-> Text Code)
  [["" 0 0] (#.Identifier "" name)])

(def: (nat\encode value)
  (-> Nat Text)
  (loop [input value
         output ""]
    (let [digit (case (n.% 10 input)
                  0 "0"
                  1 "1"
                  2 "2"
                  3 "3"
                  4 "4"
                  5 "5"
                  6 "6"
                  7 "7"
                  8 "8"
                  9 "9"
                  _ (undefined))
          output' ("lux text concat" digit output)
          input' (n./ 10 input)]
      (if (n.= 0 input')
        output'
        (recur input' output')))))

(macro: .public (zipped tokens state)
  {#.doc (doc "Create list zippers with the specified number of input lists."
              (def: .public zipped/2 (zipped 2))
              (def: .public zipped/3 (zipped 3))
              (zipped/3 xs ys zs)
              ((zipped 3) xs ys zs))}
  (case tokens
    (^ (list [_ (#.Nat num_lists)]))
    (if (n.> 0 num_lists)
      (let [(^open ".") ..functor
            indices (..indices num_lists)
            type_vars (: (List Code) (map (|>> nat\encode identifier$) indices))
            zipped_type (` (All [(~+ type_vars)]
                             (-> (~+ (map (: (-> Code Code) (function (_ var) (` (List (~ var)))))
                                          type_vars))
                                 (List [(~+ type_vars)]))))
            vars+lists (|> indices
                           (map inc)
                           (map (function (_ idx)
                                  (let [base (nat\encode idx)]
                                    [(identifier$ base)
                                     (identifier$ ("lux text concat" base "'"))]))))
            pattern (` [(~+ (map (function (_ [v vs]) (` (#.Item (~ v) (~ vs))))
                                 vars+lists))])
            g!step (identifier$ "0step0")
            g!blank (identifier$ "0,0")
            list_vars (map product.right vars+lists)
            code (` (: (~ zipped_type)
                       (function ((~ g!step) (~+ list_vars))
                         (case [(~+ list_vars)]
                           (~ pattern)
                           (#.Item [(~+ (map product.left vars+lists))]
                                   ((~ g!step) (~+ list_vars)))

                           (~ g!blank)
                           #.End))))]
        (#.Right [state (list code)]))
      (#.Left "Cannot zipped 0 lists."))

    _
    (#.Left "Wrong syntax for zipped")))

(def: .public zipped/2 (zipped 2))
(def: .public zipped/3 (zipped 3))

(macro: .public (zipped_with tokens state)
  {#.doc (doc "Create list zippers with the specified number of input lists."
              (def: .public zipped_with/2 (zipped_with 2))
              (def: .public zipped_with/3 (zipped_with 3))
              (zipped_with/2 + xs ys)
              ((zipped_with 2) + xs ys))}
  (case tokens
    (^ (list [_ (#.Nat num_lists)]))
    (if (n.> 0 num_lists)
      (let [(^open ".") ..functor
            indices (..indices num_lists)
            g!return_type (identifier$ "0return_type0")
            g!func (identifier$ "0func0")
            type_vars (: (List Code) (map (|>> nat\encode identifier$) indices))
            zipped_type (` (All [(~+ type_vars) (~ g!return_type)]
                             (-> (-> (~+ type_vars) (~ g!return_type))
                                 (~+ (map (: (-> Code Code) (function (_ var) (` (List (~ var)))))
                                          type_vars))
                                 (List (~ g!return_type)))))
            vars+lists (|> indices
                           (map inc)
                           (map (function (_ idx)
                                  (let [base (nat\encode idx)]
                                    [(identifier$ base)
                                     (identifier$ ("lux text concat" base "'"))]))))
            pattern (` [(~+ (map (function (_ [v vs]) (` (#.Item (~ v) (~ vs))))
                                 vars+lists))])
            g!step (identifier$ "0step0")
            g!blank (identifier$ "0,0")
            list_vars (map product.right vars+lists)
            code (` (: (~ zipped_type)
                       (function ((~ g!step) (~ g!func) (~+ list_vars))
                         (case [(~+ list_vars)]
                           (~ pattern)
                           (#.Item ((~ g!func) (~+ (map product.left vars+lists)))
                                   ((~ g!step) (~ g!func) (~+ list_vars)))

                           (~ g!blank)
                           #.End))))]
        (#.Right [state (list code)]))
      (#.Left "Cannot zipped_with 0 lists."))

    _
    (#.Left "Wrong syntax for zipped_with")))

(def: .public zipped_with/2 (zipped_with 2))
(def: .public zipped_with/3 (zipped_with 3))

(def: .public (last xs)
  (All [a] (-> (List a) (Maybe a)))
  (case xs
    #.End
    #.None

    (#.Item x #.End)
    (#.Some x)
    
    (#.Item x xs')
    (last xs')))

(def: .public (inits xs)
  {#.doc (doc "For a list of size N, yields the first N-1 elements."
              "Will yield a #.None for empty lists.")}
  (All [a] (-> (List a) (Maybe (List a))))
  (case xs
    #.End
    #.None

    (#.Item x #.End)
    (#.Some #.End)
    
    (#.Item x xs')
    (case (inits xs')
      #.None
      (undefined)

      (#.Some tail)
      (#.Some (#.Item x tail)))
    ))

(def: .public concat
  {#.doc (doc "The sequential combination of all the lists.")}
  (All [a] (-> (List (List a)) (List a)))
  (\ ..monad join))

(implementation: .public (with monad)
  {#.doc (doc "Enhances a monad with List functionality.")}
  (All [M] (-> (Monad M) (Monad (All [a] (M (List a))))))

  (def: &functor (functor.compose (get@ #monad.&functor monad) ..functor))

  (def: in (|>> (\ ..monad in) (\ monad in)))
  
  (def: (join MlMla)
    (do {! monad}
      [lMla MlMla
       ... TODO: Remove this version ASAP and use one below.
       lla (for {@.old
                 (: ((:parameter 0) (List (List (:parameter 1))))
                    (monad.seq ! lMla))}
                (monad.seq ! lMla))]
      (in (concat lla)))))

(def: .public (lift monad)
  {#.doc (doc "Wraps a monadic value with List machinery.")}
  (All [M a] (-> (Monad M) (-> (M a) (M (List a)))))
  (\ monad map (\ ..monad in)))

(def: .public (enumeration xs)
  {#.doc "Pairs every element in the list with its index, starting at 0."}
  (All [a] (-> (List a) (List [Nat a])))
  (loop [idx 0
         xs xs]
    (case xs
      #.End
      #.End

      (#.Item x xs')
      (#.Item [idx x] (recur (inc idx) xs')))))

(macro: .public (when tokens state)
  {#.doc (doc "Can be used as a guard in (co)monadic be/do expressions."
              (do monad
                [value (do_something 1 2 3)
                 ..when (passes_test? value)]
                (do_something_else 4 5 6)))}
  (case tokens
    (^ (.list test then))
    (#.Right [state (.list (` (.if (~ test)
                                (~ then)
                                (.list))))])

    _
    (#.Left "Wrong syntax for when")))