#lang scheme/base

(require (planet chongkai/sml/ml-package)
         (prefix-in s: srfi/43)
         scheme/match
         (only-in (planet chongkai/sml/ml-primitives)
                  SOME? SOME SOME-content
                  NONE? NONE
                  Size? Size
                  Subscript? Subscript))

(provide Vector-struct)

(define-package Vector-struct
  (maxLen fromList tabulate length sub extract concat mapi map appi app foldli foldri foldl foldr)
  
  (define maxLen +inf.0)
  
  (define (fromList l)
    (vector->immutable-vector
     (list->vector l)))
  
  (define tabulate
    (match-lambda
      ((vector n f)
       (call-with-exception-handler
        (lambda (e)
          (if (exn:break? e)
              e
          (Size (current-continuation-marks))))
        (lambda ()
          (vector->immutable-vector
           (build-vector n f)))))))
  
  (define sub
    (match-lambda
      ((vector v i)
       (call-with-exception-handler
        (lambda (e)
          (if (exn:break? e)
              e
          (Subscript (current-continuation-marks))))
        (lambda () (vector-ref v i))))))
  
  (define extract
    (match-lambda
      ((vector v s (? NONE?))
       (call-with-exception-handler
        (lambda (e)
          (if (exn:break? e)
              e
          (Subscript (current-continuation-marks))))
        (lambda ()
          (let* ((l (vector-length v))
                 (n (make-vector (- v s))))
            (let lp ((i s))
              (when (< i l)
                (vector-set! n (- i s)
                             (vector-ref v i))
                (lp (add1 i))))
            (vector->immutable-vector n)))))
      ((vector v s (? SOME? (app SOME-content e)))
       (call-with-exception-handler
        (lambda (e)
          (if (exn:break? e)
              e
          (Subscript (current-continuation-marks))))
        (lambda ()
          (let ((n (make-vector (- e s))))
            (let lp ((i s))
              (when (< i e)
                (vector-set! n (- i s)
                             (vector-ref v i))
                (lp (add1 i))))
            (vector->immutable-vector n)))))))
  
  (define ((mapi f) s)
    (let* ((stop (vector-length s))
           (newvec (make-vector stop)))
      (let lp ((j 0))
        (when (< j stop)
          (vector-set! newvec j (f (vector j (vector-ref s j))))
          (lp (add1 j))))
      (vector->immutable-vector newvec)))
  
  (define ((map f) vec)
    (let* ((len (vector-length vec))
           (new-vec (make-vector len)))
      (do ((i 0 (add1 i)))
        ((= i len) (vector->immutable-vector new-vec))
        (vector-set! new-vec i
                     (f (vector-ref vec i))))))
  
  (define ((appi f) s)
    (let ((stop (vector-length s)))
      (let lp ((i 0))
        (when (< i stop)
          (f (vector-immutable i (vector-ref s i)))
          (lp (add1 i))))))
  
  (define ((app f) v)
    (for-each f (vector->list v))
    (vector-immutable))
  
  (define (((foldli f) init) s)
    (let ((stop (vector-length s)))
      (let lp ((i 0)
               (acc init))
        (if (< i stop)
            (lp (add1 i)
                (f (vector-immutable i (vector-ref s i) acc)))
            acc))))
  
  (define (((foldri f) init) s)
    (let lp ((j (sub1 (vector-length s)))
             (acc init))
      (if (>= j 0)
          (lp (sub1 j)
              (f (vector-immutable j (vector-ref s j) acc)))
          acc)))
  
  (define (((foldl f) init) v)
    (let ((end (vector-length v)))
      (let lp ((i 0)
               (acc init))
        (if (< i end)
            (lp (add1 i)
                (f (vector-immutable (vector-ref v i) acc)))
            acc))))
  
  (define (((foldr f) init) v)
    (let lp ((i (sub1 (vector-length v)))
             (acc init))
      (if (negative? i)
          acc
          (lp (sub1 i)
              (f (vector-immutable (vector-ref v i) acc))))))
  
  (define length vector-length)
  (define concat s:vector-concatenate))