open Ir

let rec optimise expr =
  match expr with
  | Map { f; p; loc } ->
      let p' = optimise p in
      (match p' with
       | Pure { value; _ } ->
           Pure { loc; value = [%expr [%e f] [%e value]] }
       | Map { f = f2; p = p2; _ } ->
           Map { loc; f = [%expr fun x -> [%e f] ([%e f2] x)]; p = p2 }
       | _ -> Map { loc; f; p = p' })

  | SeqRight { p; q; loc } ->
      let p' = optimise p in
      let q' = optimise q in
      let p' = match p' with
        | Many { p = inner; loc = l } -> SkipMany { loc = l; p = inner }
        | Some_ { p = inner; loc = l } -> SkipSome { loc = l; p = inner }
        | _ -> p'
      in
      (match p' with
       | Pure _ -> q'
       | _ -> SeqRight { loc; p = p'; q = q' })

  | SeqLeft { p; q; loc } ->
      let p' = optimise p in
      let q' = optimise q in
      let q' = match q' with
        | Many { p = inner; loc = l } -> SkipMany { loc = l; p = inner }
        | Some_ { p = inner; loc = l } -> SkipSome { loc = l; p = inner }
        | _ -> q'
      in
      (match q' with
       | Pure _ -> p'
       | _ -> SeqLeft { loc; p = p'; q = q' })

  | Apply { pf; px; loc } ->
      let pf' = optimise pf in
      let px' = optimise px in
      (match pf', px' with
       | Pure { value = f; _ }, Pure { value = x; _ } ->
           Pure { loc; value = [%expr [%e f] [%e x]] }
       | Pure { value = f; _ }, _ ->
           Map { loc; f; p = px' }
       | _, Pure { value = x; _ } ->
           Map { loc; f = [%expr fun f -> f [%e x]]; p = pf' }
       | _ -> Apply { loc; pf = pf'; px = px' })

  | Alt { p; q; loc } ->
      let p' = optimise p in
      let q' = optimise q in
      (match p' with
       | Fail _ -> q'
       | _ -> Alt { loc; p = p'; q = q' })

  | Bind { p; f; loc } ->
      let p' = optimise p in
      (match p' with
       | Pure { value; _ } ->
           Bind { loc; p = Pure { loc; value = [%expr ()] }; f = [%expr fun () -> [%e f] [%e value]] }
       | _ -> Bind { loc; p = p'; f })

  | Attempt { p; loc } ->
      Attempt { loc; p = optimise p }
  | Lookahead { p; loc } ->
      Lookahead { loc; p = optimise p }
  | NotFollowedBy { p; loc } ->
      NotFollowedBy { loc; p = optimise p }
  | Many { p; loc } ->
      Many { loc; p = optimise p }
  | Some_ { p; loc } ->
      Some_ { loc; p = optimise p }
  | SkipMany { p; loc } ->
      SkipMany { loc; p = optimise p }
  | SkipSome { p; loc } ->
      SkipSome { loc; p = optimise p }
  | TakeWhile _ -> expr
  | Optional { p; loc } ->
      Optional { loc; p = optimise p }
  | Option { default; p; loc } ->
      Option { loc; default; p = optimise p }
  | SepBy { p; sep; loc } ->
      SepBy { loc; p = optimise p; sep = optimise sep }
  | SepBy1 { p; sep; loc } ->
      SepBy1 { loc; p = optimise p; sep = optimise sep }
  | EndBy { p; sep; loc } ->
      EndBy { loc; p = optimise p; sep = optimise sep }
  | EndBy1 { p; sep; loc } ->
      EndBy1 { loc; p = optimise p; sep = optimise sep }
  | ManyTill { p; end_; loc } ->
      ManyTill { loc; p = optimise p; end_ = optimise end_ }
  | Count { n; p; loc } ->
      Count { loc; n; p = optimise p }
  | Between { open_; close; p; loc } ->
      Between { loc; open_ = optimise open_; close = optimise close; p = optimise p }
  | ChainL { p; op; loc } ->
      ChainL { loc; p = optimise p; op = optimise op }
  | ChainL1 { p; op; loc } ->
      ChainL1 { loc; p = optimise p; op = optimise op }
  | ChainR { p; op; loc } ->
      ChainR { loc; p = optimise p; op = optimise op }
  | ChainR1 { p; op; loc } ->
      ChainR1 { loc; p = optimise p; op = optimise op }
  | Choice { ps; label; loc } ->
      let ps' = List.map optimise ps in
      let ps'' = List.filter (function Fail _ -> false | _ -> true) ps' in
      Choice { loc; ps = (if ps'' = [] then ps' else ps''); label }
  | Label { p; label; loc } ->
      Label { loc; p = optimise p; label }
  | Memo { name; p; loc } ->
      Memo { loc; name; p = optimise p }

  | Pure _ | Fail _ | Satisfy _ | Any _ | Eof _ | Token _ | Tokens _
  | OneOf _ | NoneOf _ | Cut _ | Fix _ | Var _ -> expr