sieve-of-atkin.scm

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; SECTION ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Macros
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;; Unicode is the future
(define-macro (λ . args)
  `(lambda ,@args))

(define-macro (cons-lazy a b)
  `(cons ,a (delay ,b)))

;; Imperative lazy foreach (only allows a single variable, not multiple like
;; Racket's)
(define-macro (for-lazy var . rest)
  (let ([name      (car  var)]
        [list-expr (cadr var)])
    `(let ([--list-- ,list-expr])
       (let --loop-- ([--list-- --list--])
         (if (null? --list--)
             (void)
             (let ([,name (car --list--)])
               ,@rest
               (--loop-- (cdr-lazy --list--))))))))

(define-macro (list* first . rest)
  (if (null? rest)
      first
      `(cons ,first (list* ,@rest))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; SECTION ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Functional programming helpers
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (curry f . curried-args)
  (lambda args (apply f (append curried-args args))))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; SECTION ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Lazy helpers
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (cdr-lazy lst) (force (cdr lst)))

(define (foldl-lazy f default lst)
  (if (null? lst)
      default
      (let ([fst (car lst)]
            [snd (cdr-lazy lst)])
        (f fst (foldl-lazy f default snd)))))

(define (map-lazy f lst)
  (if (null? lst)
      '()
      (let ([fst (car lst)]
            [snd (cdr-lazy lst)])
        (cons-lazy (f fst) (map-lazy f snd)))))

(define (filter-lazy pred? lst)
  (if (null? lst)
      '()
      (let ([fst (car lst)]
            [snd (cdr-lazy lst)])
        (if (pred? fst)
            (cons-lazy fst (filter-lazy pred? snd))
            (filter-lazy pred? snd)))))

(define (range-lazy lo hi)
  (range-step-lazy lo hi (if (< lo hi) 1 -1)))

(define (range-step-lazy lo hi step)
  (let ([op (if (< step 0) <= >=)])
    (range-step-op-lazy lo hi step op)))

(define (range-step-op-lazy lo hi step op)
  (if (op lo hi)
      '()
      (cons-lazy
       lo
       (range-step-op-lazy (+ lo step)
                           hi
                           step
                           op))))

(define (append-lazy fst . lsts)
  (if (null? lsts)
      fst
      (if (null? fst)
          (apply append-lazy lsts)
          (cons-lazy (car fst)
                     (apply append-lazy
                            (force (cdr fst))
                            lsts)))))

(define (cross-product-lazy fst snd)
  (define (inner lst elem)
    (map-lazy (curry list elem)
              lst))
  (apply-append-lazy (map-lazy (curry inner snd)
                               fst)))

(define (take-while-lazy pred? lst)
  (if (null? lst)
      lst
      (let ([fst (car lst)]
            [snd (cdr-lazy lst)])
        (if (pred? fst)
            (cons-lazy fst (take-while-lazy pred? snd))
            '()))))

(define (apply-append-lazy lsts)
  (let ([fst  (car lsts)]
        [rest (cdr-lazy lsts)])
    (if (null? rest)
        fst
        (if (null? fst)
            (apply-append-lazy rest)
            (cons-lazy (car fst)
                       (apply-append-lazy
                        (cons-lazy (cdr-lazy fst)
                                   rest)))))))

(define (skip-while-lazy pred? lst)
  (if (null? lst)
      lst
      (if (pred? (car lst))
          (skip-while-lazy pred? (cdr-lazy lst))
          lst)))

(define (list->lazy-list lst)
  (if (null? lst)
      '()
      (cons-lazy (car lst) (list->lazy-list (cdr lst)))))

(define (lazy-list->list lst)
  (foldl-lazy cons '() lst))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; SECTION ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Sieve of Atkin-related helpers
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (solutions-4x^2+y^2 limit)
  (define search-space
    (cross-product-lazy (range-lazy 1      ; All x's
                                    (sqrt (/ (- limit 1) 4)))
                        (range-step-lazy 1 ; Odd y's
                                         (sqrt (- limit 4))
                                         2)))
  (filter-lazy
   (curry >= limit)
   (map-lazy
    (curry apply
           (λ (x y) (+ (* 4 (square x))
                       (square y))))
    search-space)))

(define (solutions-3x^2+y^2 limit)
  (define search-space
    (cross-product-lazy (range-step-lazy 1 ; Odd x's
                                         (sqrt (/ (- limit 1) 3))
                                         2)
                        (range-step-lazy 2 ; Even y's
                                         (sqrt (- limit 3))
                                         2)))
  (filter-lazy
   (curry >= limit)
   (map-lazy
    (curry apply
           (λ (x y) (+ (* 3 (square x))
                       (square y))))
    search-space)))

(define (solutions-3x^2-y^2 limit)
  ;; Solves quadratic 3x² - (x - 1)² = limit (i.e. the x value that gives the
  ;; maximum possible value of 3x² - y²)
  (define max-x (/ (- (sqrt (- (* 2 limit) 1)) 1) 2))
  (define search-space
    (apply-append-lazy
     (map-lazy (λ (x)
                 (map-lazy (curry list x)
                           (range-step-lazy (- x 1) 0 -2)))
               (range-lazy 2 max-x))))
  (filter-lazy
   (curry >= limit)
   (map-lazy
    (curry apply
           (λ (x y) (- (* 3 (square x))
                       (square y))))
    search-space)))

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; SECTION ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;; Sieve of Atkin, see https://en.wikipedia.org/wiki/Sieve_of_Atkin
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(define (primes-below n)
  ;; For performance reasons we use constant-time indexable `vector' instead of
  ;; list, and set it mutably instead of purely.
  ;; We waste space for 0, 1, 2, 3, etc. (which we will never need to check),
  ;; but it makes the implementation more readable.
  (define possible-solutions (make-vector n #f))
  (define is-prime? (curry vector-ref possible-solutions))
  (define set-prime! (curry vector-set! possible-solutions))
  (define (flip! n) (set-prime! n (not (is-prime? n))))
  (define wheel-hits '(1 7 11 13 17 19 23 29 31 37 41 43 47 49 53 59))
  (define check-nums
    (take-while-lazy
     (curry >= n)
     (map-lazy
      (λ (pair) (+ (* 60 (car pair))
                   (cadr pair)))
      (cross-product-lazy
       (range-lazy 0 (/ n 60))
       (list->lazy-list wheel-hits)))))
  (for-lazy [i (solutions-4x^2+y^2 n)]
            (case (remainder i 60)
              [(1 13 17 29 37 41 49 53)
               (flip! i)]))
  (for-lazy [i (solutions-3x^2+y^2 n)]
            (case (remainder i 60)
              [(7 19 31 43)
               (flip! i)]))
  (for-lazy [i (solutions-3x^2-y^2 n)]
            (case (remainder i 60)
              [(11 23 47 59)
               (flip! i)]))
  (for-lazy [i (filter-lazy
                (λ (x)
                  (and (>= x 7)
                       (is-prime? x)))
                check-nums)]
            (define not-prime
              (take-while-lazy (curry >= n)
                               (map-lazy (curry * i i)
                                         check-nums)))
            (for-lazy [j not-prime]
                      (set-prime! j #f)))
  (list* 2 3 5
         (lazy-list->list
          (filter-lazy is-prime? (range-lazy 7 n)))))