path: root/stdlib/source/lux/concurrency/stm.lux

(;module:
  lux
  (lux (control functor
                applicative
                monad)
       [io #- run]
       (data (coll [list #* "List/" Functor<List> Fold<List>]
                   [dict #+ Dict]
                   ["Q" queue])
             [product]
             [text]
             maybe
             [number "Nat/" Codec<Text,Nat>]
             text/format)
       [macro]
       (macro [code]
              ["s" syntax #+ syntax: Syntax])
       (concurrency [atom #+ Atom atom]
                    ["P" promise #+ Promise "Promise/" Monad<Promise>]
                    [frp])
       ))

(type: (Var-State a)
  {#value a
   #observers (Dict Text (-> a (IO Unit)))})

(type: #export (Var a)
  {#;doc "A mutable cell containing a value, and observers that will be alerted of any change to it."}
  (Atom (Var-State a)))

(type: (Tx-Frame a)
  {#var (Var a)
   #original a
   #current a})

(type: Tx
  (List (Ex [a] (Tx-Frame a))))

(type: #export (STM a)
  {#;doc "A computation which updates a transaction and produces a value."}
  (-> Tx [Tx a]))

(def: #export (var value)
  {#;doc "Creates a new STM var, with a default value."}
  (All [a] (-> a (Var a)))
  (atom;atom {#value value
              #observers (dict;new text;Hash<Text>)}))

(def: raw-read
  (All [a] (-> (Var a) a))
  (|>. atom;get io;run (get@ #value)))

(def: (find-var-value var tx)
  (All [a] (-> (Var a) Tx (Maybe a)))
  (|> tx
      (find (function [[_var _original _current]]
              (is (:! (Var Unit) var)
                  (:! (Var Unit) _var))))
      (:: Monad<Maybe> map (function [[_var _original _current]]
                             _current))
      (:! (Maybe ($ +0)))))

(def: #export (read var)
  (All [a] (-> (Var a) (STM a)))
  (function [tx]
    (case (find-var-value var tx)
      (#;Some value)
      [tx value]

      #;None
      (let [value (raw-read var)]
        [(#;Cons [var value value] tx)
         value]))))

(def: #export (read! var)
  {#;doc "Reads var immediately, without going through a transaction."}
  (All [a] (-> (Var a) (IO a)))
  (|> var
      atom;get
      (:: Functor<IO> map (get@ #value))))

(def: (update-tx-value var value tx)
  (All [a] (-> (Var a) a Tx Tx))
  (case tx
    #;Nil
    #;Nil
    
    (#;Cons [_var _original _current] tx')
    (if (is (:! (Var ($ +0)) var)
            (:! (Var ($ +0)) _var))
      (#;Cons [(:! (Var ($ +0)) _var)
               (:! ($ +0) _original)
               (:! ($ +0) value)]
              tx')
      (#;Cons [_var _original _current]
              (update-tx-value var value tx')))
    ))

(def: #export (write value var)
  (All [a] (-> a (Var a) (STM Unit)))
  (function [tx]
    (case (find-var-value var tx)
      (#;Some _)
      [(update-tx-value var value tx)
       []]

      #;None
      [(#;Cons [var (raw-read var) value] tx)
       []])))

(def: #export (write! new-value var)
  {#;doc "Writes value to var immediately, without going through a transaction."}
  (All [a] (-> a (Var a) (IO Unit)))
  (do Monad<IO>
    [old (atom;get var)
     #let [old-value (get@ #value old)
           new (set@ #value new-value old)]
     succeeded? (atom;compare-and-swap old new var)]
    (if succeeded?
      (do @
        [_ (|> old
               (get@ #observers)
               dict;values
               (mapM @ (function [f] (f new-value))))]
        (wrap []))
      (write! new-value var))))

(def: #export (follow target)
  {#;doc "Creates a channel that will receive all changes to the value of the given var."}
  (All [a] (-> (Var a) (IO (frp;Chan a))))
  (let [head (frp;chan ($ +0))
        chan-var (var head)
        observer (function [label value]
                   (case (io;run (|> chan-var raw-read (frp;write value)))
                     #;None
                     ## By closing the output Chan, the
                     ## observer becomes obsolete.
                     (atom;update (function [[value observers]]
                                    [value (dict;remove label observers)])
                                  target)

                     (#;Some tail')
                     (write! tail' chan-var)))]
    (do Monad<IO>
      [_ (atom;update (function [[value observers]]
                        (let [label (Nat/encode (List/fold (function [key base]
                                                             (case (Nat/decode key)
                                                               (#;Left _)
                                                               base

                                                               (#;Right key-num)
                                                               (n.max key-num base)))
                                                           +0
                                                           (dict;keys observers)))]
                          [value (dict;put label (observer label) observers)]))
                      target)]
      (wrap head))))

(struct: #export _ (Functor STM)
  (def: (map f fa)
    (function [tx]
      (let [[tx' a] (fa tx)]
        [tx' (f a)]))))

(struct: #export _ (Applicative STM)
  (def: functor Functor<STM>)

  (def: (wrap a)
    (function [tx] [tx a]))

  (def: (apply ff fa)
    (function [tx]
      (let [[tx' f] (ff tx)
            [tx'' a] (fa tx')]
        [tx'' (f a)]))))

(struct: #export _ (Monad STM)
  (def: applicative Applicative<STM>)

  (def: (join mma)
    (function [tx]
      (let [[tx' ma] (mma tx)]
        (ma tx')))))

(def: #export (update! f var)
  {#;doc "Will update a Var's value, and return a tuple with the old and the new values."}
  (All [a] (-> (-> a a) (Var a) (IO [a a])))
  (io (loop [_ []]
        (let [(^@ state [value observers]) (io;run (atom;get var))
              value' (f value)]
          (if (io;run (atom;compare-and-swap state
                                             [value' observers]
                                             var))
            [value value']
            (recur []))))))

(def: #export (update f var)
  {#;doc "Will update a Var's value, and return a tuple with the old and the new values."}
  (All [a] (-> (-> a a) (Var a) (STM [a a])))
  (do Monad<STM>
    [a (read var)
     #let [a' (f a)]
     _ (write a' var)]
    (wrap [a a'])))

(def: (can-commit? tx)
  (-> Tx Bool)
  (every? (function [[_var _original _current]]
            (is _original (raw-read _var)))
          tx))

(def: (commit-var [_var _original _current])
  (-> (Ex [a] (Tx-Frame a)) Unit)
  (if (is _original _current)
    []
    (io;run (write! _current _var))))

(def: fresh-tx Tx (list))

(def: pending-commits
  (Var (Ex [a] [(STM a) (Promise a)]))
  (var (:!! [])))

(def: commit-processor-flag
  (Atom Bool)
  (atom false))

(def: (process-commit commits)
  (-> (frp;Chan [(STM Unit) (Promise Unit)])
      (Promise Unit))
  (do P;Monad<Promise>
    [?head+tail commits]
    (case ?head+tail
      (#;Cons [stm-proc output] tail)
      (do @
        [#let [[finished-tx value] (stm-proc fresh-tx)]]
        (exec (if (can-commit? finished-tx)
                (exec (List/map commit-var finished-tx)
                  (io;run (P;resolve value output))
                  [])
                (exec (io;run (write! [stm-proc output] pending-commits))
                  []))
          (process-commit tail)))

      #;Nil
      (undefined)
      )))

(def: init-processor!
  (IO Unit)
  (do Monad<IO>
    [flag (atom;get commit-processor-flag)]
    (if flag
      (wrap [])
      (do @
        [was-first? (atom;compare-and-swap flag true commit-processor-flag)]
        (if was-first?
          (do Monad<IO>
            [inputs (follow pending-commits)]
            (exec (process-commit (:! (frp;Chan [(STM Unit) (Promise Unit)])
                                      inputs))
              (wrap [])))
          (wrap [])))
      )))

(def: #export (commit stm-proc)
  {#;doc "Commits a transaction and returns its result (asynchronously).

          Note that a transaction may be re-run an indeterminate number of times if other transactions involving the same variables successfully commit first.

          For this reason, it's important to note that transactions must be free from side-effects, such as I/O."}
  (All [a] (-> (STM a) (Promise a)))
  (let [output (P;promise ($ +0))]
    (exec (io;run init-processor!)
      (io;run (write! [stm-proc output] pending-commits))
      output)))