(module modular-arithmetic
  (xgcd mod+ mod- mod* mod/ modexpt with-modulus)
  (import
    scheme
    (chicken base)
    (chicken bitwise)
    srfi-1)
  (import-for-syntax
    scheme
    (chicken base)
    matchable)

;; Extended GCD
(define (xgcd a b)
  (let-values ([(q m) (quotient&modulo a b)])
    (if (zero? m)
	(values 0 1)
	(let-values ([(x y) (xgcd b m)])
	  (values y (- x (* y q)))))))

;; Modular addition generator
(define ((mod+ modulus) . ns)
  (fold (lambda (n a) (modulo (+ a n) modulus)) 0 ns))

;; Modular subtraction and negation generator
(define (mod- modulus)
  (case-lambda
    [(a)
     (modulo (- a) modulus)]
    [(a . ns)
     (fold (lambda (n a) (modulo (- a n) modulus)) a ns)]))

;; Modular multiplication generator
(define ((mod* modulus) . ns)
  (fold (lambda (n a) (modulo (* a n) modulus)) 1 ns))

;; Modular division and inversion generator
(define (mod/ modulus)
  (define (inverse a)
    (if (= a 1)
	1
	(let-values ([(1/a n) (xgcd a modulus)])
	  (if (zero? n)
	      (error 'mod/ "operand and modulus are not coprime" a modulus)
	      1/a))))
  (case-lambda
    [(a)
     (modulo (inverse a) modulus)]
    [(a . ns)
     (fold (lambda (n a) (modulo (* a (inverse n)) modulus)) a ns)]))

;; Modular exponentiation generator
(define (modexpt modulus)
  (let ([* (mod* modulus)]
        [/ (mod/ modulus)])
    (lambda (base exponent)
      (let loop ([a 1]
                 [base (if (negative? exponent) (/ base) base)]
                 [exponent (abs exponent)])
        (if (positive? exponent)
	    (loop (if (zero? (bitwise-and exponent 1)) a (* a base))
                  (* base base)
                  (arithmetic-shift exponent -1))
            a)))))

;; Syntax to overload +, add1, -, sub1, *, / and expt with modular versions
(define-syntax with-modulus
  (er-macro-transformer
   (lambda (stx rename id=)
     (match stx
       [(with-modulus modulus body ...)
        (let ([~let (rename 'let)]
	      [~letrec (rename 'letrec)]
	      [~modulus (rename 'modulus)])
          `(,~let ([,~modulus ,modulus])
             (,~letrec ([+ (,(rename 'mod+) ,~modulus)]
			[add1 (lambda (n) (+ n 1))]
			[- (,(rename 'mod-) ,~modulus)]
			[sub1 (lambda (n) (- n 1))]
			[* (,(rename 'mod*) ,~modulus)]
			[/ (,(rename 'mod/) ,~modulus)]
			[expt (,(rename 'modexpt) ,~modulus)])
               ,@body)))]))))

)

;; vim: set ai et ts=8 sts=2 sw=2 ft=scheme: ;;