(.module:
  [lux #*
   [control
    [monad (#+ do)]]
   [data
    ["." error]]
   ["." function]
   [type
    ["tc" check]]
   ["." macro]]
  [// (#+ Operation)
   ["/." //
    ["." extension]]])

(def: #export (with-type expected)
  (All [a] (-> Type (Operation a) (Operation a)))
  (extension.localized (get@ #.expected) (set@ #.expected)
                       (function.constant (#.Some expected))))

(def: #export (with-env action)
  (All [a] (-> (tc.Check a) (Operation a)))
  (function (_ (^@ stateE [bundle state]))
    (case (action (get@ #.type-context state))
      (#error.Error error)
      ((///.fail error) stateE)

      (#error.Success [context' output])
      (#error.Success [[bundle (set@ #.type-context context' state)]
                       output]))))

(def: #export with-fresh-env
  (All [a] (-> (Operation a) (Operation a)))
  (extension.localized (get@ #.type-context) (set@ #.type-context)
                       (function.constant tc.fresh-context)))

(def: #export (infer actualT)
  (-> Type (Operation Any))
  (do ///.Monad<Operation>
    [expectedT (extension.lift macro.expected-type)]
    (with-env
      (tc.check expectedT actualT))))

(def: #export (with-inference action)
  (All [a] (-> (Operation a) (Operation [Type a])))
  (do ///.Monad<Operation>
    [[_ varT] (..with-env
                tc.var)
     output (with-type varT
              action)
     knownT (..with-env
              (tc.clean varT))]
    (wrap [knownT output])))