advent-of-code-2022/day-12/array-helpers.scm

(library (array-helpers)
  (export array-len-in-dim
          arrays->hash-table
          array-display
          array-map

          array-next-index
          array-cell-ref-vec
          array-index-of
          array-indices-of)
  (import
   (except (rnrs base)
           let-values
           map
           error
           vector-map)
   (only (guile)
         lambda* λ
         current-output-port
         ;; arrays
         array-shape
         array-ref
         array-cell-ref
         array-map!
         array-rank
         ;; display
         display
         simple-format
         current-output-port)
   ;; lists
   (srfi srfi-1)
   ;; vectors
   (srfi srfi-43)
   ;; hash tables
   (srfi srfi-69)
   (ice-9 arrays))


  (define array-len-in-dim
    (λ (arr dim)
      (let* ([shape (array-shape arr)]
             [dim-min-max (list-ref shape dim)])
        (+ (- (second dim-min-max)
              (first dim-min-max))
           1))))


  (define arrays->hash-table
    (lambda* (keys-arr vals-arrs #:optional (equal-func equal?))
      (let ([rows (array-len-in-dim keys-arr 0)]
            [cols (array-len-in-dim keys-arr 1)]
            [table (make-hash-table equal-func)])
        (let iter-rows ([row-ind 0])
          (let iter-cols ([col-ind 0])
            (cond
             [(< row-ind rows)
              (cond
               [(< col-ind cols)
                (hash-table-set! table
                                 (array-ref keys-arr row-ind col-ind)
                                 (array-ref vals-arrs row-ind col-ind))
                (iter-cols (+ col-ind 1))]
               [else (iter-rows (+ row-ind 1))])]
             [else table]))))))


  (define array-display
    (lambda* (landscape
              #:optional (port (current-output-port))
              #:key (formatter (λ (elem) elem)))
      (let ([rows (array-len-in-dim landscape 0)]
            [cols (array-len-in-dim landscape 1)])
        (let iter-rows ([row-ind 0])
          (let iter-cols ([col-ind 0])
            (cond
             [(>= row-ind rows) 'done]
             [(>= col-ind cols)
              (display "\n" (current-output-port))
              (iter-rows (+ row-ind 1))]
             [else
              (display (formatter (array-cell-ref landscape row-ind col-ind)) port)
              (iter-cols (+ col-ind 1))]))))))


  (define array-map
    (λ (proc src-arr)
      (define target-arr (array-copy src-arr))
      (array-map! target-arr proc src-arr)
      target-arr))


  (define array-next-index
    (λ (shape-vec indices-vec max-dim)
      "Increment one of indices in INDICES-VEC for which the
following conditions are true:

1. There is no index in INDICES-VEC that sits at a later
position (higher index) than MAX-DIM of INDICES-VEC and is
not yet at its maximum. The maximum is specified by the
SHAPE-VEC. Each index in INDICES-VEC has a corresponding
minimum and maximum in SHAPE-VEC at the same position.

After incrementing the index in INDICES-VEC, all later
indices (at a higher index of INDICES-VEC) are set to their
corresponding minimum to get a correct indices vector."
      (cond
       [(>= max-dim 0)
        (let ([index (vector-ref indices-vec max-dim)]
              [max-for-index (second (vector-ref shape-vec max-dim))])
          (cond
           [(< index max-for-index)
            ;; Copy the vector to not mutate argument.
            (let ([updated-indices-vec (vector-copy indices-vec)])
              ;; Increase index at position.
              (vector-set! updated-indices-vec
                           max-dim
                           (+ (vector-ref updated-indices-vec max-dim) 1))
              ;; Set later indices to their corresponding minimum.
              (let ([indices-vec-len (vector-length indices-vec)])
                (let iter ([dim° (+ max-dim 1)])
                  (cond
                   [(< dim° indices-vec-len)
                    (let ([minimum-ind-val (first (vector-ref shape-vec dim°))])
                      (vector-set! updated-indices-vec dim° minimum-ind-val)
                      (iter (+ dim° 1)))]
                   [else updated-indices-vec]))))]
           [else
            ;; Increment next higher dimension.
            (array-next-index shape-vec
                              indices-vec
                              (- max-dim 1))]))]
       [else #f])))


  (define array-cell-ref-vec
    (λ (arr indices-vec)
      "array-cell-ref takes a variable number of arguments,
depending on how many dimensions the array has (what its
rank is). However, apply works with lists as last arguments,
not with vectors. To avoid having to convert vectors of
indices to lists of indices every time such a vector of
indices is used with an array, array-cell-ref-vec can be
used."
      (let ([vec-len (vector-length indices-vec)])
        (let iter ([index-into-indices-vec° 0] [result arr])
          (cond
           [(< index-into-indices-vec° vec-len)
            (let ([cell-index (vector-ref indices-vec index-into-indices-vec°)])
              (iter (+ index-into-indices-vec° 1)
                    (array-cell-ref result cell-index)))]
           [else result])))))


  (define array-index-of
    (lambda* (arr pred #:optional (start-indices #f))
      "Return the index of the first element in ARR which
satisfies the predicate PRED."
      (let* ([shape (array-shape arr)]
             [shape-vec (list->vector shape)]
             [rank (array-rank arr)]
             [initial-indices
              (if start-indices
                  start-indices
                  (vector-map (λ (_i elem) (car elem)) shape-vec))])
        (let iter ([indices° initial-indices])
          (cond
           [indices°
            (cond
             ;; If the array element satisfies the
             ;; predicate, return the indices of the
             ;; element.
             [(pred (array-cell-ref-vec arr indices°))
              indices°]
             [else
              ;; Potential optimization: Make better use of
              ;; the rank argument, so that array-next-index
              ;; does not have to search unnecessarily for
              ;; the index to increment.
              (iter (array-next-index shape-vec indices° (- rank 1)))])]
           ;; No index found at which the predicate would be
           ;; satisfied.
           [else #f])))))


  (define array-indices-of
    (λ (arr pred)
      (let* ([shape (array-shape arr)]
             [shape-vec (list->vector shape)]
             [rank (array-rank arr)]
             [initial-indices (vector-map (λ (_i elem) (car elem)) shape-vec)])
        (let iter ([indices° initial-indices])
          (cond
           [indices°
            (cond
             [(pred (array-cell-ref-vec arr indices°))
              (cons indices°
                    (iter (array-next-index shape-vec indices° (- rank 1))))]
             [else (iter (array-next-index shape-vec indices° (- rank 1)))])]
           [else '()]))))))