Make the example bot cleaner

1 files changed, 12 insertions, 7 deletions

diff --git a/bin/bot.ml b/bin/bot.ml
index cb27458..197efb4 100644
--- a/bin/bot.ml
+++ b/bin/bot.ml
@@ -3,24 +3,29 @@ open Core
 open Disml
 
 let check_command (msg:Message.t) =
-    if String.is_prefix ~prefix:"!ping" msg.content then
+    let cmd, rest = match String.split ~on:' ' msg.content with
+    | hd::tl -> hd, tl
+    | [] -> "", []
+    in match cmd with
+    | "!ping" ->
         Message.reply msg "Pong!" >>> ignore
-    else if String.is_prefix ~prefix:"!spam" msg.content then
-        let count = String.chop_prefix_exn ~prefix:"!spam" msg.content |> String.strip |> Int.of_string in
+    | "spam" ->
+        let count = Option.((List.hd rest >>| Int.of_string) |> value ~default:0) in
         List.range 0 count
         |> List.iter ~f:(fun i -> Message.reply msg (string_of_int i) >>> ignore)
-    else if String.is_prefix ~prefix:"!list" msg.content then
-        let count = String.chop_prefix_exn ~prefix:"!list" msg.content |> String.strip |> Int.of_string in
+    | "!list" ->
+        let count = Option.((List.hd rest >>| Int.of_string) |> value ~default:0) in
         let list = List.range 0 count
         |> List.sexp_of_t Int.sexp_of_t
         |> Sexp.to_string_hum in
         Message.reply msg list >>> ignore
-    else if String.is_prefix ~prefix:"!fold" msg.content then
-        let count = String.chop_prefix_exn ~prefix:"!fold" msg.content |> String.strip |> Int.of_string in
+    | "!fold" ->
+        let count = Option.((List.hd rest >>| Int.of_string) |> value ~default:0) in
         let list = List.range 0 count
         |> List.fold ~init:0 ~f:(+)
         |> Int.to_string in
         Message.reply msg list >>> ignore
+    | _ -> ()
 
 let main () =
     let token = match Sys.getenv "DISCORD_TOKEN" with