;;; CHICKEN Transducers - Transducers for working with foldable data types.
;;;
;;; Copyright (c) 2022 Jeremy Steward
;;;
;;; Permission is hereby granted, free of charge, to any person obtaining a copy
;;; of this software and associated documentation files (the "Software"), to deal
;;; in the Software without restriction, including without limitation the rights
;;; to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
;;; copies of the Software, and to permit persons to whom the Software is
;;; furnished to do so, subject to the following conditions:
;;;
;;; The above copyright notice and this permission notice shall be included in all
;;; copies or substantial portions of the Software.
;;;
;;; THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
;;; IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
;;; FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
;;; AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
;;; LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
;;; OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
;;; SOFTWARE.

(module (transducers base) ()
  (import (except (scheme) map for-each))
  (import (chicken base)
          (chicken module)
          (chicken fixnum)
          (only type-checks-numbers
                check-non-negative-fixnum
                check-positive-fixnum))

  ;; Reduced values record members
  (export <reduced>
          make-reduced
          reduced?
          unwrap
          preserving-reduced)

  ;; Collectors
  (export collect-any
          collect-all
          collect-count
          collect-max
          collect-min
          collect-first
          collect-last
          collect-sum
          collect-product
          collect-void)

  ;; Transduce procedures
  (export transduce for-each)

  ;; Transducers
  (export map
          filter
          remove
          drop
          drop-while
          take
          take-while
          chunks
          chunks-exact
          enumerate
          inspect)

  ;; Syntax
  (export define-chain-transducer
          define-flatten-transducer
          define-interleave-transducer
          define-zip-transducer)

  ;; A record type expressing the concept of a reduced value - i.e. a value
  ;; that tells the transduction to stop due to a completed step in the
  ;; transducer.
  ;;
  ;; This is needed to be separate from "exhausted," which represents a
  ;; collection that has no more remaining values.
  (define-record-type <reduced>
    ;; Constructor for creating a `reduced` value.
    (make-reduced value)
    ;; Predicate to check if something is a reduced value.
    reduced?
    ;; Unwraps a reduced value and returns the value.
    (value unwrap))

  ;; A reducer that aggregates all items as if applying `(or (pred? <>) ...)` across
  ;; all items transduced over.
  (define (collect-any pred? #!optional (sentinel #f))
    (case-lambda
      (() sentinel)
      ((result) result)
      ((result item)
       (let ((test (pred? item)))
         (if (or result test)
           (make-reduced test)
           #f)))))

  ;; A reducer that aggregates all items as if applying `(and (pred? <>) ...)`
  ;; across all items transduced over.
  (define (collect-all pred? #!optional (sentinel #t))
    (case-lambda
      (() sentinel)
      ((result) result)
      ((result item)
       (let ((test (pred? item)))
         (if (and result test)
           test
           (make-reduced #f))))))

  ;; A reducer that aggregates the total number of items transduced over.
  (define (collect-count #!optional (sentinel 0))
    (case-lambda
      (() sentinel)
      ((result) result)
      ((result item)
       (fx+ result 1))))

  ;; A reducer that aggregates the maximum of the items transduced over.
  (define (collect-max #!optional (sentinel -inf.0))
    (case-lambda
      (() sentinel)
      ((result) result)
      ((result item)
       (max result item))))

  ;; A reducer that aggregates the minimum of the items transduced over.
  (define (collect-min #!optional (sentinel +inf.0))
    (case-lambda
      (() sentinel)
      ((result) result)
      ((result item)
       (min result item))))

  ;; A reducer that returns just the first item, or #f if there are no items.
  (define (collect-first #!optional (sentinel #f))
    (case-lambda
      (() sentinel)
      ((result) result)
      ((result item) (make-reduced item))))

  ;; A reducer that returns just the last item, or #f if there are no items.
  (define (collect-last #!optional (sentinel #f))
    (case-lambda
      (() sentinel)
      ((result) result)
      ((result item) item)))

  ;; A reducer that returns the sum of the items using `+`
  (define (collect-sum #!optional (sentinel 0))
    (case-lambda
      (() sentinel)
      ((result) result)
      ((result item) (+ result item))))

  ;; A reducer that returns the product of the items using `*`
  (define (collect-product #!optional (sentinel 1))
    (case-lambda
      (() sentinel)
      ((result) result)
      ((result item) (* result item))))

  ;; A reducer that collects all items but always returns an undefined result
  ;;
  ;; Primarily useful for implementing for-each, which is just a transducer
  ;; that operates over all items in an iterable but then returns an undefined
  ;; result at the end.
  (define (collect-void #!optional (sentinel (void)))
    (case-lambda
      (() (void))
      ((result) result)
      ((result item) result)))

  ;; A procedure that transduces over an iterable with the provided xform,
  ;; collector, and iterable.
  (define (transduce folder transducer collector data)
    (let* ((xf (transducer collector))
           (result (folder xf (collector) data)))
      (xf result)))

  ;; Transducer that performs a mapping operation over each item in the iterable.
  ;;
  ;; For every item that passes through this transducer, (f item) is passed on
  ;; to the reducing function.
  (define (map f)
    (lambda (reducer)
      (case-lambda
        (() (reducer))
        ((result) (reducer result))
        ((result item)
         (reducer result (f item))))))

  ;; Transducer that performs a filtering operation over each item in the iterable.
  ;;
  ;; For every item that passes through this transducer, the item is passed
  ;; forward iff it satisfies `pred`
  (define (filter pred)
    (lambda (reducer)
      (case-lambda
        (() (reducer))
        ((result) (reducer result))
        ((result item)
         (if (pred item)
           (reducer result item)
           result)))))

  ;; Transducer that operates like filter, but will only forward items iff
  ;; the item satisfies `(compose not pred)`.
  (define (remove pred)
    (filter (compose not pred)))

  ;; A transducer that drops `n` number of items before forwarding all
  ;; remaining items.
  (define (drop n)
    (check-non-negative-fixnum 'drop n 'n)
    (lambda (reducer)
      (let ((n n))
        (case-lambda
          (() (reducer))
          ((result) (reducer result))
          ((result item)
           (set! n (fx- n 1))
           (if (fx>= n 0)
             result
             (reducer result item)))))))

  ;; A transducer that drops items while pred is true. After the first instance
  ;; of pred returning false, continues to forward items along to the reducer.
  (define (drop-while pred)
    (lambda (reducer)
      (let ((continue-to-drop #t))
        (case-lambda
          (() (reducer))
          ((result) (reducer result))
          ((result item)
           (if continue-to-drop
             (if (pred item)
               result
               (begin
                 (set! continue-to-drop #f)
                 (reducer result item)))
             (reducer result item)))))))

  ;; A transducer that takes `n` number of items before dropping all remaining
  ;; items.
  (define (take n)
    (check-positive-fixnum 'take n 'n)
    (lambda (reducer)
      (let ((n n))
        (case-lambda
          (() (reducer))
          ((result) (reducer result))
          ((result item)
           (set! n (fx- n 1))
           (if (fx> n 0)
             (reducer result item)
             (make-reduced
               (reducer result item))))))))

  ;; A transducer that takes items as long as `pred` is true, and then drops
  ;; all remaining items after `pred` first turns false.
  (define (take-while pred)
    (lambda (reducer)
      (case-lambda
        (() (reducer))
        ((result) (reducer result))
        ((result item)
         (if (pred item)
           (reducer result item)
           (make-reduced result))))))

  ;; Helper function that wraps a reduced value twice as traversing full
  ;; iterables will unwrap them. `flatten` is a good example: it re-uses its
  ;; reducer on the item using `traverse-iterable.` If that reduction finishes
  ;; early and returns a reduced value, then it will be immediately unwrapped
  ;; and we will want to keep moving forward with the transduction.
  (define (preserving-reduced reducer)
    (lambda (acc x)
      (let ((ret (reducer acc x)))
        (if (reduced? ret)
          (make-reduced ret)
          ret))))

  ;; Defines a transducer that flattens every item that satisfies `type?` by
  ;; folding over that type.
  (define-syntax define-flatten-transducer
    (syntax-rules ()
      ((_ name type? type-fold)
       (define name
         (lambda (reducer)
           (case-lambda
             (() (reducer))
             ((result) (reducer result))
             ((result item)
              (if (type? item)
                (type-fold (preserving-reduced (name reducer)) result item)
                (reducer result item)))))))))

  ;; Chains the iterable into the transduction, as if appending one iterable to
  ;; another.
  (define-syntax define-chain-transducer
    (syntax-rules ()
      ((_ name type-fold)
       (define (name iterable)
         (lambda (reducer)
           (case-lambda
             (() (reducer))
             ((result)
              (reducer (type-fold reducer result iterable)))
             ((result item)
              (reducer result item))))))))

  ;; Collects items into a collection of `n` items defined by `collector`.
  ;;
  ;; Collector can be any valid collector, e.g. collect-list, every, any, +, *,
  ;; etc.
  ;;
  ;; NOTE: chunks will collect up-to `n` items in each collection. If the end
  ;; of a collection is reached before `n` items are collected, then the last
  ;; chunk will only have as many items (less than `n`) that were collected.
  (define (chunks n collector)
    (check-positive-fixnum 'segment n 'n)
    (lambda (reducer)
      (let ((i 0)
            (collection (collector)))
        (case-lambda
          (() (reducer))
          ((result)
           ;; If there is anything remaining in our collector when finalizing we
           ;; collect here. This last collection will have `i` number of items.
           (if (fx> i 0)
             (let ((new-result (reducer result (collector collection))))
               (if (reduced? new-result)
                 (reducer (unwrap new-result))
                 (reducer new-result)))
             (reducer result)))
          ((result item)
           (let ((new-i (fx+ i 1))
                 (new-collection (collector collection item)))
             (if (fx< new-i n)
               (begin
                 (set! collection new-collection)
                 (set! i new-i)
                 result)
               (begin
                 (set! collection (collector))
                 (set! i 0)
                 (reducer result (collector new-collection))))))))))

  ;; Collects items into a collection of `n` items defined by `collector`.
  ;;
  ;; Collector can be any valid collector, e.g. collect-list, every, any, +, *,
  ;; etc.
  ;;
  ;; NOTE: Unlike chunks, if the underlying transduction is exhausted prior to
  ;; collecting `n` chunks then the remaining items are discarded and only
  ;; chunks of exactly `n` length will be returned.
  (define (chunks-exact n collector)
    (check-positive-fixnum 'segment n 'n)
    (lambda (reducer)
      (let ((i 0)
            (collection (collector)))
        (case-lambda
          (() (reducer))
          ((result) (reducer result))
          ((result item)
           (let ((new-i (fx+ i 1))
                 (new-collection (collector collection item)))
             (if (fx< new-i n)
               (begin
                 (set! collection new-collection)
                 (set! i new-i)
                 result)
               (begin
                 (set! collection (collector))
                 (set! i 0)
                 (reducer result (collector new-collection))))))))))

  ;; Creates a transducer that can interleave elements of a given collection
  ;; type into the transduced items.
  (define-syntax define-interleave-transducer
    (syntax-rules ()
      ((_ name make-state done? next update)
       (define (name collection)
         (lambda (reducer)
           (let ((state (make-state collection)))
             (case-lambda
               (() (reducer))
               ((result) (reducer result))
               ((result item)
                (if (done? collection state)
                  (make-reduced result)
                  (let ((interleave-item (next collection state))
                        (next-result (reducer result item)))
                    (set! state (update collection state))
                    (if (reduced? next-result)
                      next-result
                      (reducer next-result interleave-item))))))))))))

  ;; Creates a transducer that can interleave elements of a given collection
  ;; type into the transduced items.
  (define-syntax define-zip-transducer
    (syntax-rules ()
      ((_ name make-state done? next update)
       (define (name collection)
         (lambda (reducer)
           (let ((state (make-state collection)))
             (case-lambda
               (() (reducer))
               ((result) (reducer result))
               ((result item)
                (if (done? collection state)
                  (make-reduced result)
                  (let ((zip-item (next collection state)))
                    (set! state (update collection state))
                    (reducer result (cons item zip-item))))))))))))

  ;; Enumerates the items in a transducer. Pairs the reducer with an index such
  ;; that for every item in the reducer you get a pair that has a zero-based
  ;; index describing the ordering of the item.
  ;;
  ;; E.g. if you enumerate the list '(a b c) you'll get the output '(0 . a)
  ;; followed by '(1 . b) and '(2 . c).
  (define enumerate
    (lambda (reducer)
      (let ((n 0))
        (case-lambda
          (() (reducer))
          ((result) (reducer result))
          ((result item)
           (let ((enumerated-item (cons n item)))
             (set! n (fx+ n 1))
             (reducer result enumerated-item)))))))

  ;; Inspect an item with some function f and pass that item along.
  ;;
  ;; Unlike map, this isn't meant to actually modify the item at all, but
  ;; instead to just run some function `f` on it and then forward that item
  ;; through the transduction. Especially useful for debugging or performing
  ;; side-effects.
  ;;
  ;; It is an error to modify the item passed to `f` in any way.
  (define (inspect f)
    (lambda (reducer)
      (case-lambda
        (() (reducer))
        ((result) (reducer result))
        ((result item)
         (f item)
         (reducer result item)))))

  ;; for-each is a shorthand form for writing transducers that return an
  ;; unspecified / undefined value.
  ;;
  ;; Mostly this is useful when using side-effects / mutation, but can also be
  ;; a short-hand for running programs that aren't meant to return anything in
  ;; the first place.
  (define (for-each folder transducer data)
    (transduce folder transducer (collect-void) data))

  ;; End-of-module
  )