;;; Copyright (c) 2023 Jeremy Steward
;;; All rights reserved.
;;;
;;; Redistribution and use in source and binary forms, with or without
;;; modification, are permitted provided that the following conditions are met:
;;;
;;; 1. Redistributions of source code must retain the above copyright notice,
;;; this list of conditions and the following disclaimer.
;;;
;;; 2. Redistributions in binary form must reproduce the above copyright notice,
;;; this list of conditions and the following disclaimer in the documentation
;;; and/or other materials provided with the distribution.
;;;
;;; 3. Neither the name of the copyright holder nor the names of its
;;; contributors may be used to endorse or promote products derived from this
;;; software without specific prior written permission.
;;;
;;; THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
;;; AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
;;; IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
;;; ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
;;; LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
;;; CONSEQUENTIAL DAMAGES  INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
;;; SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
;;; INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
;;; CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
;;; ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
;;; POSSIBILITY OF SUCH DAMAGE.

;;; Transducers, procedures for working with transducers, and a transducible
;;; type-class for arrays.

;; A procedure which folds over the indices of the array in lexicographic
;; order, applying the procedure `(f collection index)` over each index, with
;; `sentinel` standing in as the initial `collection` for the first element.
;;
;; This is more or less a convenience for folding over the shape of the array,
;; but without needing to individually extract the shape.
(define-checked (array-interval-fold f sentinel (array array?))
  (let ((shape (array-shape array)))
    (interval-fold f sentinel shape)))

;; A procedure which folds over the indices of the array in reverse
;; lexicographic order, applying the procedure `(f collection index)` over each
;; index, with `sentinel` standing in as the initial `collection` for the first
;; element.
;;
;; This is analagous to array-interval-fold but folded over in reverse.
(define-checked (reverse-array-interval-fold f sentinel (array array?))
  (let ((shape (array-shape array)))
    (reverse-interval-fold f sentinel shape)))

;; A procedure which folds over the elements of the array in lexicographic
;; order, applying the procedure `(f collection element)` over each element,
;; with `sentinel` standing in as the initial `collection` for the first
;; element.
;;
;; Lexicographic ordering here means that it follows the order of the indices.
;; If you were to access each individual index using array-index-fold and then
;; do (array-ref array idx) on each index, you would access the correct
;; elements in the correct order.
(define-checked (array-fold f sentinel (array array?))
  (let* ((storage-class (array-storage-class array))
         (storage-object (array-storage-object array))
         (storage-ref (storage-class-ref storage-class))
         (stride (array-stride array))
         (rank (vector-length stride)))
    (interval-fold ((map
                      (lambda (idx)
                        (storage-ref storage-object (index->offset idx stride rank))))
                    f)
                   sentinel
                   (array-interval array))))

;; A procedure which folds over the elements of the array in reverse
;; lexicographic order, applying the procedure `(f collection element)` over
;; each element, with `sentinel` standing in as the initial `collection` for
;; the first element.
;;
;; Reverse-lexicographic ordering here means that it follows the order of the
;; indices, but in reverse. This is analagous to array-fold but operates over
;; the elements in the reverse order of the indices.
(define-checked (reverse-array-fold f sentinel (array array?))
  (let* ((storage-class (array-storage-class array))
         (storage-object (array-storage-object array))
         (storage-ref (storage-class-ref storage-class))
         (stride (array-stride array))
         (rank (vector-length stride)))
    (reverse-interval-fold ((map
                              (lambda (idx)
                                (storage-ref storage-object (index->offset idx stride rank))))
                            f)
                           sentinel
                           (array-interval array))))

(define-flatten-transducer flatten-array array? array-fold)
(define-chain-transducer chain-array array-fold)

(define-flatten-transducer reverse-flatten-array array? reverse-array-fold)
(define-chain-transducer reverse-chain-array reverse-array-fold)

(define-checked (zip-array (array array?))
  (let* ((storage-class (array-storage-class array))
         (storage-object (array-storage-object array))
         (storage-ref (storage-class-ref storage-class))
         (stride (array-stride array))
         (rank (vector-length stride)))
    (compose
      (zip-interval (array-interval array))
      (map
        (lambda (pair)
          (let ((value (car pair))
                (index (cdr pair)))
            (cons value
                  (storage-ref storage-object (index->offset index stride rank)))))))))

(define-checked (reverse-zip-array (array array?))
  (let* ((storage-class (array-storage-class array))
         (storage-object (array-storage-object array))
         (storage-ref (storage-class-ref storage-class))
         (stride (array-stride array))
         (rank (vector-length stride)))
    (compose
      (reverse-zip-interval (array-interval array))
      (map
        (lambda (pair)
          (let ((value (car pair))
                (index (cdr pair)))
            (cons value
                  (storage-ref storage-object (index->offset index stride rank)))))))))

(define flatten-pair
  (lambda (reducer)
    (case-lambda
      (() (reducer))
      ((result) (reducer result))
      ((result item)
       (if (pair? item)
         (let ((result (reducer result (car item))))
           (if (reduced? result)
             result
             (reducer result (cdr item))))
         (reducer result item))))))

(define-checked (interleave-array (array array?))
  (compose (zip-array array) flatten-pair))

(define-checked (reverse-interleave-array (array array?))
  (compose (reverse-zip-array array) flatten-pair))

(define-checked (collect-array (storage-class storage-class?) (dimension vector?))
  (let* ((rank (vector-length dimension))
         (len (transduce fixnum-range-fold
                         (map (lambda (i) (vector-ref dimension i)))
                         (collect-product)
                         (range 0 rank)))
         (collector ((transducible-collector (storage-class-transducible storage-class)) len)))
    (unless (fx>? rank 0)
      (error 'collect-array
             "Collecting into an array requires a dimension with a rank of at least 1."
             dimension
             rank))
    (case-lambda
      (() (collector))
      ((result)
       (let ((storage-object (collector result)))
         (make-array-from-storage storage-class dimension storage-object)))
      ((result item)
       (collector result item)))))

(define array-transducible
  (make-transducible array-fold
                     collect-array
                     flatten-array
                     chain-array
                     interleave-array
                     zip-array))

(define reverse-array-transducible
  (make-transducible reverse-array-fold
                     collect-array
                     reverse-flatten-array
                     reverse-chain-array
                     reverse-interleave-array
                     reverse-zip-array))

;; Returns a newly allocated array (with a newly allocated shape of the
;; specified index-class and storage object of the specified storage-class)
;; that has the specified dimension `dim`. The values of the elements are
;; computed by calling proc on every possible index of the array in
;; lexicographic order.
(define-checked (array-tabulate (storage-class storage-class?) (dimension vector?) proc)
  (transduce interval-fold
             (map proc)
             (collect-array storage-class dimension)
             (make-default-interval dimension)))