;; SPDX-License-Identifier: EUPL-1.2
;; SPDX-FileCopyrightText: 2025 Uko Kokņevičs <perkontevs@gmail.com>
(defpackage :ukkoclot/src/rw-lock
  (:documentation "Implementation of a shared/read-write lock.")
  (:use :c2cl :bt2 :iterate)
  (:import-from :alexandria :whichever :with-gensyms)
  (:import-from :named-readtables :in-readtable)
  ;; These imports make this SBCL-dependent, but technically speaking we could live without them.
  (:import-from :sb-sys
   :allow-with-interrupts :with-local-interrupts :without-interrupts)
  (:import-from :serapeum :->)
  (:import-from :ukkoclot/src/readtable :readtable)
  (:export
   #:rw-lock
   #:rw-lock-p
   #:make-rw-lock
   #:acquire-read-lock
   #:release-read-lock
   #:with-read-lock
   #:acquire-write-lock
   #:release-write-lock
   #:with-write-lock))
(in-package :ukkoclot/src/rw-lock)
(in-readtable readtable)

;; TODO: Use atomic-integer in best-case for read locks to decrease contention
(defstruct (rw-lock (:constructor make-rw-lock%))
  (lock nil :type lock :read-only t)
  (reader-cv nil :type condition-variable :read-only t)
  (writer-cv nil :type condition-variable :read-only t)
  (waiting-readers 0 :type (integer 0))
  (waiting-writers 0 :type (integer 0))
  (active-readers 0 :type (integer 0))
  (active-writer nil :type boolean))

(defvar *counter-lock* (make-lock :name "rw-lock::*counter-lock*"))

(defvar *counter* 0)

(-> gen-name () string)
(defun gen-name ()
  "Generate a name for a rw-lock"
  (format nil "Read-Write Lock ~A"
          (with-lock-held (*counter-lock*)
            (incf *counter*))))

(-> make-rw-lock (&key (:name string)) rw-lock)
(defun make-rw-lock (&key (name (gen-name)))
  (check-type name string)
  (make-rw-lock%
   :lock (make-lock :name #f"{name}'s internal lock")
   :reader-cv (make-condition-variable :name #f"{name}'s internal reader-cv")
   :writer-cv (make-condition-variable :name #f"{name}'s internal writer-cv")))

(-> wakeup-waiters (rw-lock) (values &optional))
(defun wakeup-waiters (rw-lock)
  ;; NOTE: RW-LOCK-LOCK HAS TO BE TAKEN!
  (declare (type rw-lock rw-lock))
  (with-slots (reader-cv writer-cv waiting-readers waiting-writers) rw-lock
    (cond
      ((zerop waiting-readers) (condition-notify writer-cv))
      ((zerop waiting-writers) (condition-broadcast reader-cv))
      (t (whichever (condition-notify writer-cv)
                    (condition-broadcast reader-cv)))))
  (values))

(-> acquire-read-lock (rw-lock &key (:wait boolean)) boolean)
(defun acquire-read-lock (rw-lock &key (wait t))
  ;; TODO: timeout
  (check-type rw-lock rw-lock)
  (with-slots (lock reader-cv active-readers active-writer waiting-readers) rw-lock
    (let ((lock-acquired))
      (without-interrupts
        (unwind-protect
             (when (setf lock-acquired
                         (allow-with-interrupts
                           (acquire-lock lock :wait wait)))
               (incf waiting-readers)
               (with-local-interrupts
                 (when (if wait
                           (iter
                             (unless active-writer
                               (return t))
                             (condition-wait reader-cv lock))
                           (not active-writer))
                   (incf active-readers)
                   t)))
          (when lock-acquired
            (decf waiting-readers)
            (release-lock lock)))))))

(-> release-read-lock (rw-lock) rw-lock)
(defun release-read-lock (rw-lock)
  (check-type rw-lock rw-lock)
  (with-slots (lock active-readers active-writer) rw-lock
    (with-lock-held (lock)
      (assert (> active-readers 0) nil "Trying to release a read lock that's not taken!")
      (decf active-readers)
      (when (zerop active-readers)
        (wakeup-waiters rw-lock))))
  rw-lock)

(defmacro with-read-lock ((rw-lock) &body body)
  ;; TODO: timeout
  (with-gensyms (lock-acquired lock-value)
    `(let ((,lock-acquired)
           (,lock-value ,rw-lock))
       (without-interrupts
         (unwind-protect
              (when (setf ,lock-acquired
                          (allow-with-interrupts
                           (acquire-read-lock ,lock-value)))
                (with-local-interrupts ,@body))
           (when ,lock-acquired
             (release-read-lock ,lock-value)))))))

(-> acquire-write-lock (rw-lock &key (:wait boolean)) boolean)
(defun acquire-write-lock (rw-lock &key (wait t))
  ;; TODO: timeout
  (check-type rw-lock rw-lock)
  (with-slots (lock writer-cv active-readers active-writer waiting-writers) rw-lock
    (let ((lock-acquired))
      (without-interrupts
        (unwind-protect
             (when (setf lock-acquired
                         (allow-with-interrupts
                           (acquire-lock lock :wait wait)))
               (incf waiting-writers)
               (with-local-interrupts
                 (when (if wait
                           (iter
                             (unless (or (> active-readers 0) active-writer)
                               (return t))
                             (condition-wait writer-cv lock))
                           (and (zerop active-readers) (not active-writer)))
                   (setf active-writer t)
                   t)))
          (when lock-acquired
            (decf waiting-writers)
            (release-lock lock)))))))

(-> release-write-lock (rw-lock) rw-lock)
(defun release-write-lock (rw-lock)
  (check-type rw-lock rw-lock)
  (with-slots (lock active-readers active-writer) rw-lock
    (with-lock-held (lock)
      (assert active-writer nil "Trying to release a write lock that's not taken!")
      (setf active-writer nil)
      (wakeup-waiters rw-lock)))
  rw-lock)

(defmacro with-write-lock ((rw-lock) &body body)
  ;; TODO: timeout
  (with-gensyms (lock-acquired lock-value)
    `(let ((,lock-acquired)
           (,lock-value ,rw-lock))
       (without-interrupts
         (unwind-protect
              (when (setf ,lock-acquired
                          (allow-with-interrupts
                            (acquire-write-lock ,lock-value)))
                (with-local-interrupts ,@body))
           (when ,lock-acquired
             (release-write-lock ,lock-value)))))))