open Ir

type first =
  | Epsilon
  | Predicate of Ppxlib.expression
  | Token of Ppxlib.expression
  | Tokens of Ppxlib.expression
  | OneOf of Ppxlib.expression
  | Any
  | Unknown

type t = {
  firsts : first list;
  nullable : bool;
}

let empty = { firsts = []; nullable = false }
let epsilon = { firsts = [Epsilon]; nullable = true }
let any = { firsts = [Any]; nullable = false }
let unknown = { firsts = [Unknown]; nullable = false }

let of_predicate e = { firsts = [Predicate e]; nullable = false }
let of_token e = { firsts = [Token e]; nullable = false }
let of_tokens e = { firsts = [Tokens e]; nullable = false }
let of_one_of e = { firsts = [OneOf e]; nullable = false }

let union a b = {
  firsts = a.firsts @ b.firsts;
  nullable = a.nullable || b.nullable;
}

let seq a b =
  if a.nullable then union { a with nullable = false } b
  else a

let rec compute env = function
  | Pure _ -> epsilon
  | Fail _ -> empty
  | Satisfy { pred; _ } -> of_predicate pred
  | Any _ -> any
  | Eof _ -> epsilon
  | Token { tok; _ } -> of_token tok
  | Tokens { toks; _ } -> of_tokens toks
  | OneOf { toks; _ } -> of_one_of toks
  | NoneOf _ -> unknown
  | Bind { p; _ } -> compute env p
  | Map { p; _ } -> compute env p
  | Apply { pf; px; _ } -> seq (compute env pf) (compute env px)
  | SeqLeft { p; q; _ } -> seq (compute env p) (compute env q)
  | SeqRight { p; q; _ } -> seq (compute env p) (compute env q)
  | Alt { p; q; _ } -> union (compute env p) (compute env q)
  | Attempt { p; _ } -> compute env p
  | Cut _ -> epsilon
  | Lookahead { p; _ } -> { (compute env p) with nullable = true }
  | NotFollowedBy _ -> epsilon
  | Many _ -> epsilon
  | Some_ { p; _ } -> compute env p
  | SkipMany _ -> epsilon
  | SkipSome { p; _ } -> compute env p
  | TakeWhile { at_least_one = false; _ } -> epsilon
  | TakeWhile { pred; at_least_one = true; _ } -> of_predicate pred
  | Optional { p; _ } -> { (compute env p) with nullable = true }
  | Option { p; _ } -> { (compute env p) with nullable = true }
  | SepBy _ -> epsilon
  | SepBy1 { p; _ } -> compute env p
  | EndBy _ -> epsilon
  | EndBy1 { p; _ } -> compute env p
  | ManyTill { p; _ } -> union (compute env p) epsilon
  | Count { p; _ } -> compute env p
  | Between { open_; _ } -> compute env open_
  | ChainL _ -> epsilon
  | ChainL1 { p; _ } -> compute env p
  | ChainR _ -> epsilon
  | ChainR1 { p; _ } -> compute env p
  | Choice { ps; _ } -> List.fold_left (fun acc p -> union acc (compute env p)) empty ps
  | Label { p; _ } -> compute env p
  | Memo { p; _ } -> compute env p
  | Fix { names; _ } ->
      List.fold_left (fun acc name ->
        match List.assoc_opt name env with
        | Some fs -> union acc fs
        | None -> acc
      ) empty names
  | Var { name; _ } ->
      (match List.assoc_opt name env with
       | Some fs -> fs
       | None -> unknown)

let has_unknown t = List.exists (function Unknown -> true | _ -> false) t.firsts

let definitely_disjoint a b =
  if has_unknown a || has_unknown b then false
  else if a.nullable && b.nullable then false
  else
    let singleton_token = function
      | [Token _] -> true
      | [OneOf _] -> true
      | _ -> false
    in
    singleton_token a.firsts && singleton_token b.firsts

let is_static = function
  | Token _ -> true
  | Tokens _ -> true
  | OneOf _ -> true
  | _ -> false

let all_static t = List.for_all is_static t.firsts