forest-fire/001/src/forest_fire.ml

   1 open Printf
   2
   3
   4 (* ------------------------------------------------------------------------- *
   5  * Constants
   6  * ------------------------------------------------------------------------- *)
   7 let default_f = 0.01  (* Probability of spontaneous ignition *)
   8 let default_p = 1.0   (* Probability of spontaneous growth *)
   9
  10 let default_x = 80
  11 let default_y = 25
  12
  13 let char_empty   = ' '
  14 let char_tree    = 'T'
  15 let char_burning = '#'
  16
  17 let ansi_color_tree    = "\027[0;32m"  (* Green *)
  18 let ansi_color_burning = "\027[1;31m"  (* Red *)
  19 let ansi_color_off     = "\027[0m"
  20
  21 let ansi_code_clear = "\027[2J"    (* Clear screen *)
  22 let ansi_code_reset = "\027[1;1H"  (* Reset cursor position *)
  23
  24
  25 (* ------------------------------------------------------------------------- *
  26  * Types
  27  * ------------------------------------------------------------------------- *)
  28 type cell_state =
  29   | Empty | Tree | Burning
  30
  31
  32 type direction =
  33   | N | NE | E | SE | S | SW | W | NW
  34
  35
  36 type options =
  37   { size : int * int
  38   ; prob : float * float
  39   }
  40
  41
  42 (* ------------------------------------------------------------------------- *
  43  * Utils
  44  * ------------------------------------------------------------------------- *)
  45
  46 (* Hack to sleep less than 1 sec *)
  47 let minisleep subsec =
  48   ignore (Unix.select [] [] [] subsec)
  49
  50
  51 let term_clear () =
  52   print_string ansi_code_clear
  53
  54
  55 let term_reset () =
  56   print_string ansi_code_reset
  57
  58
  59 let get_opts argv =
  60   let usage = ""
  61
  62   and f = ref default_f
  63   and p = ref default_p
  64   and x = ref default_x
  65   and y = ref default_y in
  66
  67   let speclist =
  68     Arg.align
  69     [ ("-f", Arg.Set_float f, " Probability of spontaneous ignition.")
  70     ; ("-p", Arg.Set_float p, " Probability of spontaneous growth.")
  71     ; ("-x", Arg.Set_int x, " Forest width.")
  72     ; ("-y", Arg.Set_int y, " Forest height.")
  73     ]
  74   in
  75
  76   Arg.parse speclist (fun _ -> ()) usage;
  77
  78   { size = !x, !y
  79   ; prob = !f, !p
  80   }
  81
  82
  83 (* ------------------------------------------------------------------------- *
  84  * Core
  85  * ------------------------------------------------------------------------- *)
  86 let directions =
  87   [N; NE; E; SE; S; SW; W; NW]
  88
  89
  90 let offset_of_direction = function
  91   (* Direction -> x, y *)
  92   | N  ->  0, -1
  93   | NE ->  1, -1
  94   | E  ->  1,  0
  95   | SE ->  1,  1
  96   | S  ->  0,  1
  97   | SW -> -1,  1
  98   | W  -> -1,  0
  99   | NW -> -1, -1
 100
 101
 102 let offsets =
 103   List.map (offset_of_direction) directions
 104
 105
 106 let is_probable = function
 107   | probability when (Random.float 1.0) <= probability -> true
 108   | _ -> false
 109
 110
 111 let init_cell_state (_, p) = function
 112   | () when is_probable p -> Tree
 113   | ()                    -> Empty
 114
 115
 116 let init_forest (x, y) prob =
 117   Array.map (Array.map (init_cell_state prob)) (Array.make_matrix y x ())
 118
 119
 120 let string_of_state = function
 121   | Empty   -> sprintf "%c" char_empty
 122   | Tree    -> sprintf "%s%c%s" ansi_color_tree char_tree ansi_color_off
 123   | Burning -> sprintf "%s%c%s" ansi_color_burning char_burning ansi_color_off
 124
 125
 126 let new_state = function
 127   | Burning, _, _                            -> Empty
 128   | Tree,    0, (f, _)    when is_probable f -> Burning
 129   | Tree,    n_burning, _ when n_burning > 0 -> Burning
 130   | Empty,   _, (_, p)    when is_probable p -> Tree
 131   | state,   _, _                            -> state
 132
 133
 134 let print_forest forest =
 135   Array.iter
 136   (
 137     fun row ->
 138       Array.iter
 139       (
 140         fun state ->
 141           print_string (string_of_state state)
 142       )
 143       row;
 144       print_newline ()
 145   )
 146   forest
 147
 148
 149 let is_onside width height (x, y) =
 150   x >= 0 && y >= 0 && x < width && y < height
 151
 152
 153 let next_generation forest (width, height) prob =
 154   Array.mapi
 155   (
 156     fun iy row ->
 157       Array.mapi
 158       (
 159         fun ix state ->
 160           let neighbors = List.map (fun (ox, oy) -> ox + ix, oy + iy) offsets in
 161           let neighbors = List.filter (is_onside width height) neighbors in
 162           let neighbor_states = List.map (fun (x, y) -> forest.(y).(x)) neighbors in
 163           let burning_states = List.filter (fun s -> s == Burning) neighbor_states in
 164           new_state (state, (List.length burning_states), prob)
 165       )
 166       row
 167   )
 168   forest
 169
 170
 171 let rec burn forest size prob =
 172   term_reset ();
 173   print_forest forest;
 174   minisleep 0.1;
 175   burn (next_generation forest size prob) size prob
 176
 177
 178 let main argv =
 179   Random.self_init ();
 180
 181   let opts = get_opts argv in
 182   let forest = init_forest opts.size opts.prob in
 183
 184   term_clear ();
 185   burn forest opts.size opts.prob
 186
 187
 188 let () = main Sys.argv