(**************************************************************************)
(*  Copyright © 2014 Mehdi Dogguy <mehdi@dogguy.org>                      *)
(*                                                                        *)
(*  This program is free software: you can redistribute it and/or modify  *)
(*  it under the terms of the GNU Affero General Public License as        *)
(*  published by the Free Software Foundation, either version 3 of the    *)
(*  License, or (at your option) any later version, with the additional   *)
(*  exemption that compiling, linking, and/or using OpenSSL is allowed.   *)
(*                                                                        *)
(*  This program is distributed in the hope that it will be useful, but   *)
(*  WITHOUT ANY WARRANTY; without even the implied warranty of            *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU     *)
(*  Affero General Public License for more details.                       *)
(*                                                                        *)
(*  You should have received a copy of the GNU Affero General Public      *)
(*  License along with this program.  If not, see                         *)
(*  <http://www.gnu.org/licenses/>.                                       *)
(**************************************************************************)

let level = ref (Domain.recommended_domain_count ())
let set_level l = if l > 0 then level := l

module T = Domainslib.Task

let get_level () = !level

let map ?level:level' f l =
  let num_domains = (match level' with None -> !level | Some l -> l) - 1 in
  if num_domains > 0 then (
    let pool = T.setup_pool ~num_domains () in
    let l =
      T.run pool (fun () ->
          let l = List.map (fun x -> T.async pool (fun () -> f x)) l in
          List.map (fun x -> T.await pool x) l)
    in
    T.teardown_pool pool;
    l)
  else List.map f l

let iter ?level:level' f l =
  let num_domains = (match level' with None -> !level | Some l -> l) - 1 in
  if num_domains > 0 then
    let pool = T.setup_pool ~num_domains () in
    let () =
      T.run pool (fun () ->
          let l = List.map (fun x -> T.async pool (fun () -> f x)) l in
          List.iter (fun x -> T.await pool x) l)
    in
    T.teardown_pool pool
  else List.iter f l
