(module lmdb
  (database-environment?
   current-database-environment
   open-database-environment close-database-environment
   copy-database-environment
   with-transaction clear-stale-transactions
   database?
   open-database close-database drop-database
   database-ref database-set! database-exists? database-delete!
   database-fold database-walk
   database->alist alist->database)

(import
  scheme
  (chicken base)
  (chicken blob)
  (chicken fixnum)
  (chicken keyword)
  (chicken condition)
  (chicken memory)
  (chicken foreign)
  (only (chicken file) create-directory)
  srfi-1)

(foreign-declare
  "#include <lmdb.h>")

(define (type-error location message . arguments)
  (abort
   (make-composite-condition
    (make-property-condition
     'exn
     'location location
     'message message
     'arguments arguments)
    (make-property-condition
     'type))))

(define-syntax define-foreign-tagged-type
  (syntax-rules ()
    [(define-foreign-tagged-type (id nonnull-id target-type)
       predicate tag)
     (begin
       (define tag
         (gensym 'id))
       (define predicate
         (cut tagged-pointer? <> tag))
       (define-foreign-type id (c-pointer target-type)
         (lambda (v)
           (and v (if (predicate v) v (type-error 'id "bad argument type" v))))
         (lambda (v)
           (and v (tag-pointer v tag))))
       (define-foreign-type nonnull-id (nonnull-c-pointer target-type)
         (lambda (v)
           (if (predicate v) v (type-error 'nonnull-id "bad argument type" v)))
         (lambda (v)
           (tag-pointer v tag))))]))

(define-syntax define-foreign-enum-argconvert
  (syntax-rules ()
    [(define-foreign-enum-argconvert (id base-type)
       (key value) ...)
     (define (id v #!optional [seed 0])
       (case v
         [(key) (fxior (foreign-value value base-type) seed)] ...
         [else #f]))]))

(define-syntax define-foreign-enum-retconvert
  (syntax-rules ()
    [(define-foreign-enum-retconvert (id base-type)
       (key value) ...)
     (define (id v)
       (select v
         [((foreign-value value base-type)) 'key] ...
         [else #f]))]))

(define-syntax define-foreign-enum-type
  (syntax-rules ()
    [(define-foreign-enum-type (id base-type)
       argconvert retconvert
       (key value) ...)
     (begin
       (define-foreign-enum-argconvert (argconvert base-type)
         (key value) ...)
       (define-foreign-enum-retconvert (retconvert base-type)
         (key value) ...)
       (define-foreign-type id base-type
         (lambda (v)
           (or (argconvert v) (type-error 'id "unknown enumeration value" v)))
         retconvert))]))

(define-syntax let-keys+flags
  (syntax-rules ()
    [(let-keys+flags loc args-expr
                     ([key key-var key-default] ...)
                     ([flag-argconvert flags-var] ...)
       expr ...)
     (let ([key-var key-default] ...
           [flags-var 0] ...)
       (let next ([args args-expr] [pos 0])
         (when (pair? args)
           (let ([v (car args)])
             (cond
              [(flag-argconvert v flags-var)
               => (lambda (v)
                    (set! flags-var v)
                    (next (cdr args) pos))]
              ...
              [else
               (let-values ([(k v next-args next-pos)
                             (if (keyword? v)
                                 (values
                                  v (if (pair? (cdr args))
                                        (cadr args)
                                        (error 'loc "illegal arguments" args))
                                  (cddr args) pos)
                                 (values
                                  pos v
                                  (cdr args) (fx+ pos 1)))])
                 (case k
                   [(key)
                    (set! key-var v)
                    (next next-args next-pos)]
                   ...
                   [else
                    (error 'loc "illegal arguments" args)]))]))))
       expr ...)]))

;; Errors

(define (check-error loc s)
  (unless (zero? s)
    (abort
     (make-composite-condition
      (make-property-condition
       'exn
       'location loc
       'message ((foreign-lambda c-string "mdb_strerror" int) s))
      (make-property-condition
       'db
       'code s)))))

(define-values (map-resized not-found)
  (values
   (foreign-value "MDB_MAP_RESIZED" int)
   (foreign-value "MDB_NOTFOUND" int)))

;; Database environments

(define-foreign-tagged-type (database-environment nonnull-database-environment "MDB_env")
  database-environment?
  tag:database-environment)

(define current-database-environment
  (make-parameter #f))

(define-foreign-enum-argconvert (database-environment-flags unsigned-int)
  [#:fixed-map "MDB_FIXEDMAP"]
  [#:no-subdirectory "MDB_NOSUBDIR"]
  [#:read-only "MDB_RDONLY"]
  [#:write-map "MDB_WRITEMAP"]
  [#:no-meta-sync "MDB_NOMETASYNC"]
  [#:no-sync "MDB_NOSYNC"]
  [#:map-async "MDB_MAPASYNC"]
  [#:no-lock "MDB_NOLOCK"]
  [#:no-read-ahead "MDB_NORDAHEAD"])

(define-values (no-tls no-subdirectory)
  (values
   (foreign-value "MDB_NOTLS" unsigned-int)
   (foreign-value "MDB_NOSUBDIR" unsigned-int)))

(define (open-database-environment path . args)
  (let-keys+flags open-database-environment args
                  ([#:mode mode #o666]
                   [#:max-databases maxdbs #f]
                   [#:max-readers maxreaders #f]
                   [#:max-size mapsize #f])
                  ([database-environment-flags flags])
    (when (zero? (fxand flags no-subdirectory))
      (create-directory path))
    (let-location ([env database-environment #f])
      (check-error
       'open-database-environment
       ((foreign-lambda
         int "mdb_env_create"
         (nonnull-c-pointer database-environment))
        (location env)))
      (handle-exceptions exn
        (begin
          (close-database-environment env)
          (abort exn))
        (when maxdbs
          (check-error
           'open-database-environment
           ((foreign-lambda
             int "mdb_env_set_maxdbs"
             nonnull-database-environment unsigned-int)
            env maxdbs)))
        (when maxreaders
          (check-error
           'open-database-environment
           ((foreign-lambda
             int "mdb_env_set_maxreaders"
             nonnull-database-environment unsigned-int)
            env maxreaders)))
        (when mapsize
          (check-error
           'open-database-environment
           ((foreign-lambda
             int "mdb_env_set_mapsize"
             nonnull-database-environment size_t)
            env mapsize)))
        (check-error
         'open-database-environment
         ((foreign-lambda
           int "mdb_env_open"
           nonnull-database-environment nonnull-c-string unsigned-int unsigned-int)
          env path (fxior flags no-tls) mode)))
      (current-database-environment env))))

(define (close-database-environment #!optional [env (current-database-environment)])
  ((foreign-lambda
    void "mdb_env_close"
    nonnull-database-environment)
   env)
  (when (equal? env (current-database-environment))
    (current-database-environment #f))
  (void))

(define-foreign-enum-argconvert (database-environment-copy-flags unsigned-int)
  [#:compact "MDB_CP_COMPACT"])

(define (copy-database-environment path . args)
  (let-keys+flags copy-database-environment args
                  ([0 env (current-database-environment)])
                  ([database-environment-copy-flags flags])
    (when (let-location ([flags unsigned-int 0])
            (check-error
             'copy-database-environment
             ((foreign-lambda
               int "mdb_env_get_flags"
               nonnull-database-environment (c-pointer unsigned-int))
              env (location flags)))
            (zero? (fxand flags no-subdirectory)))
      (create-directory path))
    (check-error
     'copy-database-environment
     ((foreign-lambda
       int "mdb_env_copy2"
       nonnull-database-environment nonnull-c-string unsigned-int)
      env path flags))))

;; Transactions

(define-foreign-tagged-type (transaction nonnull-transaction "MDB_txn")
  transaction?
  tag:transaction)

(define current-transaction
  (make-parameter #f))

(define-foreign-enum-argconvert (transaction-flags unsigned-int)
  [#:read-only "MDB_RDONLY"])

(define (with-transaction thunk . args)
  (let-keys+flags with-transaction args
                  ([0 env (current-database-environment)])
                  ([transaction-flags flags])
    (let-location ([txn transaction #f])
      (let retry ()
        (let ([status
               ((foreign-lambda
                 int "mdb_txn_begin"
                 nonnull-database-environment transaction unsigned-int
                 (nonnull-c-pointer transaction))
                env (current-transaction) flags (location txn))])
          (if (eqv? status map-resized)
              (begin
                (check-error
                 'with-transaction
                 ((foreign-lambda
                   int "mdb_env_set_mapsize"
                   nonnull-database-environment size_t)
                  env 0))
                (retry))
              (check-error 'with-transaction status))))
      (let ([complete? #f] [success? #f])
        (dynamic-wind
          (lambda ()
            (when complete?
              (error 'with-transaction "cannot re-enter transaction")))
          (lambda ()
            (parameterize ([current-transaction txn])
              (receive results (thunk)
                (set! success? #t)
                (apply values results))))
          (lambda ()
            (if success?
                (check-error
                 'with-transaction
                 ((foreign-lambda
                   int "mdb_txn_commit"
                   nonnull-transaction)
                  txn))
                ((foreign-lambda
                  void "mdb_txn_abort"
                  nonnull-transaction)
                 txn))
            (set! complete? #t)))))))

(define (clear-stale-transactions #!optional [env (current-database-environment)])
  (let-location ([dead int 0])
    (check-error
     'clear-stale-transactions
     ((foreign-lambda
       int "mdb_reader_check"
       nonnull-database-environment (c-pointer int))
      env (location dead)))
    dead))

;; Databases

(define-record-type database
  (wrap-database dbi)
  database?
  [dbi unwrap-database])

(define-foreign-type database unsigned-int
  (lambda (v)
    (if (database? v) (unwrap-database v) (type-error 'id "bad argument type" v)))
  wrap-database)

(define-foreign-enum-argconvert (database-flags unsigned-int)
  [#:reverse-key "MDB_REVERSEKEY"]
  [#:duplicate-sort "MDB_DUPSORT"]
  [#:integer-key "MDB_INTEGERKEY"]
  [#:duplicate-fixed "MDB_DUPFIXED"]
  [#:integer-duplicate "MDB_INTEGERDUP"]
  [#:reverse-duplicate "MDB_REVERSEDUP"]
  [#:create "MDB_CREATE"])

(define (%open-database name flags)
  (let-location ([dbi database (wrap-database 0)])
    (check-error
     'open-database
     ((foreign-lambda
       int "mdb_dbi_open"
       nonnull-transaction c-string unsigned-int
       (nonnull-c-pointer database))
      (current-transaction) name flags (location dbi)))
    dbi))

(define (open-database . args)
  (let-keys+flags open-database args
                  ([0 name #f])
                  ([database-flags flags])
    (%open-database name flags)))

(define (close-database dbi #!optional [env (current-database-environment)])
  ((foreign-lambda
    void "mdb_dbi_close"
    nonnull-database-environment database)
   env dbi))

(define (database-drop-flags v #!optional [seed 0])
  (case v
    [(#:delete) (fxior 1 seed)]
    [else #f]))

(define (drop-database dbi . args)
  (let-keys+flags drop-database args
                  ()
                  ([database-drop-flags flags])
    (check-error
     'drop-database
     ((foreign-lambda
       int "mdb_drop"
       nonnull-transaction database int)
      (current-transaction) dbi flags))))

(define (blob/string-size v)
  (cond
   [(blob? v) (blob-size v)]
   [(string? v) (string-length v)]
   [else 0]))

(define-foreign-enum-argconvert (database-set!-flags unsigned-int)
  [#:no-duplicate "MDB_NODUPDATA"]
  [#:no-overwrite "MDB_NOOVERWRITE"]
  [#:append "MDB_APPEND"]
  [#:append/duplicate "MDB_APPENDDUP"])

(define (%database-set! dbi key val flags)
  (check-error
   'database-set!
   ((foreign-lambda*
     int ([nonnull-transaction txn] [database dbi]
          [size_t klen] [scheme-pointer kptr]
          [size_t vlen] [scheme-pointer vptr]
          [unsigned-int flags])
     "MDB_val key, val;"
     "key.mv_size = klen; key.mv_data = kptr;"
     "val.mv_size = vlen; val.mv_data = vptr;"
     "C_return(mdb_put(txn, dbi, &key, &val, flags));")
    (current-transaction) dbi
    (blob/string-size key) key
    (blob/string-size val) val
    flags)))

(define-values (database-ref database-set!)
  (letrec ([database-ref
            (lambda (dbi key #!optional [default (cut check-error 'database-ref not-found)])
              (let-location ([vlen size_t 0] [vptr c-pointer #f])
                (let ([status
                       ((foreign-lambda*
                         int ([nonnull-transaction txn] [database dbi]
                              [size_t klen] [scheme-pointer kptr]
                              [(c-pointer size_t) vlen] [(c-pointer c-pointer) vptr])
                         "int status;"
                         "MDB_val key, val;"
                         "key.mv_size = klen; key.mv_data = kptr;"
                         "val.mv_size = 0; val.mv_data = NULL;"
                         "status = mdb_get(txn, dbi, &key, &val);"
                         "*vlen = val.mv_size; *vptr = val.mv_data;"
                         "C_return(status);")
                        (current-transaction) dbi
                        (blob/string-size key) key
                        (location vlen) (location vptr))])
                  (cond
                    [(zero? status)
                     (let ([val (make-string vlen)])
                       (move-memory! vptr val vlen)
                       val)]
                    [(eqv? status not-found)
                     (if (procedure? default) (default) default)]
                    [else
                     (check-error 'database-ref status)]))))]
           [database-set!
            (lambda (dbi key val . args)
              (let-keys+flags database-set! args
                              ()
                              ([database-set!-flags flags])
                (%database-set! dbi key val flags)))])
    (values
     (getter-with-setter database-ref database-set!)
     database-set!)))

(define (database-exists? dbi key)
  (let ([status
         ((foreign-lambda*
           int ([nonnull-transaction txn] [database dbi]
                [size_t klen] [scheme-pointer kptr])
           "MDB_val key, val;"
           "key.mv_size = klen; key.mv_data = kptr;"
           "C_return(mdb_get(txn, dbi, &key, &val));")
          (current-transaction) dbi
          (blob/string-size key) key)])
    (cond
      [(zero? status)
       #t]
      [(eqv? status not-found)
       #f]
      [else
       (check-error 'database-exists? status)])))

(define (database-delete! dbi key #!optional val)
  (check-error
   'database-set!
   ((foreign-lambda*
     int ([nonnull-transaction txn] [database dbi]
          [size_t klen] [scheme-pointer kptr]
          [size_t vlen] [scheme-pointer vptr])
     "MDB_val key, val;"
     "key.mv_size = klen; key.mv_data = kptr;"
     "val.mv_size = vlen; val.mv_data = vptr;"
     "C_return(mdb_del(txn, dbi, &key, &val));")
    (current-transaction) dbi
    (blob/string-size key) key
    (blob/string-size val) val)))

;; Cursors

(define-foreign-tagged-type (cursor nonnull-cursor "MDB_cursor")
  cursor?
  tag:cursor)

(define close-cursor
  (foreign-lambda
   void "mdb_cursor_close"
   nonnull-cursor))

(define-values (cursor-set-range cursor-first cursor-next)
  (values
   (foreign-value "MDB_SET_RANGE" int)
   (foreign-value "MDB_FIRST" int)
   (foreign-value "MDB_NEXT" int)))

(define (database-fold proc seed dbi #!key from to< to<= limit)
  (let-location ([cursor cursor #f])
    (check-error
     'database-fold
     ((foreign-lambda
       int "mdb_cursor_open"
       nonnull-transaction database (c-pointer cursor))
      (current-transaction) dbi (location cursor)))
    (let ([complete? #f])
      (dynamic-wind
        (lambda ()
          (when complete?
            (error 'database-fold "cannot re-enter cursor walk")))
        (lambda ()
          (let loop ([op (if from cursor-set-range cursor-first)]
                     [key from] [seed seed] [fuel (or limit +inf.0)])
            (if (positive? fuel)
              (let-location ([klen size_t 0] [kptr c-pointer #f]
                             [vlen size_t 0] [vptr c-pointer #f])
                (let ([status
                       ((foreign-lambda*
                         int ([nonnull-cursor cursor] [int op]
                              [size_t k0len] [scheme-pointer k0ptr]
                              [size_t k1len] [scheme-pointer k1ptr]
                              [int inclusive]
                              [(c-pointer size_t) klen] [(c-pointer c-pointer) kptr]
                              [(c-pointer size_t) vlen] [(c-pointer c-pointer) vptr])
                         "int status;"
                         "MDB_val key0, key1, val;"
                         "key0.mv_size = k0len; key0.mv_data = k0ptr;"
                         "key1.mv_size = k1len; key1.mv_data = k1ptr;"
                         "val.mv_size = 0; val.mv_data = NULL;"
                         "status = mdb_cursor_get(cursor, &key0, &val, op);"
                         "if (status == 0 && key1.mv_data != NULL) {"
                         "  if (mdb_cmp(mdb_cursor_txn(cursor), mdb_cursor_dbi(cursor), &key0, &key1) >= inclusive) status = MDB_NOTFOUND;"
                         "}"
                         "if (status == 0) {"
                         "  *klen = key0.mv_size; *kptr = key0.mv_data;"
                         "  *vlen = val.mv_size; *vptr = val.mv_data;"
                         "}"
                         "C_return(status);")
                        cursor op
                        (blob/string-size key) key
                        (blob/string-size (or to< to<=)) (or to< to<=)
                        (if to< 0 1)
                        (location klen) (location kptr)
                        (location vlen) (location vptr))])
                  (cond
                    [(zero? status)
                     (let ([key (make-string klen)]
                           [val (make-string vlen)])
                       (move-memory! kptr key klen)
                       (move-memory! vptr val vlen)
                       (loop cursor-next key (proc key val seed) (sub1 fuel)))]
                    [(eqv? status not-found)
                     seed]
                    [else
                     (check-error 'database-ref status)])))
              seed)))
        (lambda ()
          ((foreign-lambda
            void "mdb_cursor_close"
            nonnull-cursor)
           cursor)
          (set! complete? #t))))))

(define (database-walk proc dbi . args)
  (apply
   database-fold
   (lambda (key val seed) (proc key val) seed) (void)
   dbi args))

(define (database->alist dbi #!rest args #!key from to< to<= limit)
  (fold
   (lambda (key+val rest)
     (let ([val (cdr key+val)])
       (when (pair? val) (set-cdr! key+val (reverse! val))))
     (cons key+val rest))
   '()
   (database-fold
    (if (memq #:duplicate-list args)
        (lambda (key val rest)
          (if (and (pair? rest) (equal? key (caar rest)))
              (begin
                (set-cdr! (car rest) (cons val (cdar rest)))
                rest)
              (cons (cons key (list val)) rest)))
        (lambda (key val rest)
          (cons (cons key val) rest)))
    '() dbi #:from from #:to< to< #:to<= to<= #:limit limit)))

(define (alist->database alist . args)
  (let-keys+flags open-database args
                  ([0 name #f])
                  ([database-flags dbi-flags]
                   [database-set!-flags set-flags])
    (let ([dbi (%open-database name dbi-flags)])
      (handle-exceptions exn
        (begin
          (close-database dbi)
          (abort exn))
        (for-each
          (lambda (key+val)
            (let ([key (car key+val)]
                  [val (cdr key+val)])
              (if (list? val)
                  (for-each
                    (cut %database-set! dbi key <> set-flags)
                    val)
                  (%database-set! dbi key val set-flags))))
          alist))
      dbi)))

)

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