<commit_before>open Statechart_translator_predicates

module Src = Statechart_t
module Tgt = Statechart_executable
module Bitset = Statechart_bitset

module StringMap = Map.Make(struct
  type t = string
  let compare = compare
end)

module IntMap = Map.Make(struct
  type t = int
  let compare = compare
end)

type child_type = State of int
                | OnEntry of Tgt.expression array
                | OnExit of Tgt.expression array
                | Transition of (unit -> int)
                | Invoke of Tgt.invoke
                | Pass

let get_parent ancestors =
  match ancestors with
  | parent :: _ -> parent
  | _ -> 0

exception UnparsedExpression of string
let unwrap_expr expr =
  match expr with
  | Src.Expr v -> raise (UnparsedExpression v)
  | Src.ExprValue v -> Tgt.string_expr v
  | Src.ExprParsed v -> v
  | Src.ExprUnset -> Tgt.null

let map_to_array a f =
  let size = IntMap.cardinal a in
  Array.init size (fun idx ->
    f (IntMap.find idx a)
  )

let set_mut map idx value =
  map := IntMap.add idx value !map

let append_mut_array arr value =
  arr := Array.append !arr value

let get_next r =
  let idx = !r in
  r := idx + 1;
  idx

let resolve_list map l =
  let targets = List.map (fun k -> StringMap.find k map) l in
  Array.of_list targets

let maybe_find map k =
  try
    Some (IntMap.find k map)
  with
  | _ -> None

let has_history children state_map =
  let n = Array.length children in
  let rec find i =
    if i >= n then false
    else (
      let {Tgt.State.t} = IntMap.find i state_map in
      (t == `history_deep || t == `history_shallow) || find (i + 1)
    )
  in
  find 0

let get_history_completion state states =
  (* TODO *)
  [||]

let find_initial states children =
  match Array.length children with
  | 0 -> [||]
  | n ->
    let rec find i =
      if i >= n then [|Array.get children 0|]
      else (
        let idx = Array.get children i in
        match IntMap.find idx states with
        | {Tgt.State.t=`initial} -> [|idx|]
        | _ -> find (i + 1)
      )
    in
    find 0

let get_completion state completions state_map states descendants =
  match state with
  (* TODO filter history children *)
  | {Tgt.State.t=`history_deep; idx; children} ->
    Array.append children (Bitset.to_idx_array (Array.get descendants idx))
  | {Tgt.State.t=`history_shallow; children} -> children
  | {Tgt.State.t=`parallel; children} -> children
  | {Tgt.State.idx=idx} -> (
    match maybe_find completions idx with
    | Some initial when (initial != []) -> resolve_list state_map initial
    | _ -> find_initial states state.Tgt.State.children
  )

let flatten document =
  let states = ref IntMap.empty in
  let state_idx = ref 0 in
  let state_map = ref StringMap.empty in
  let state_completions = ref IntMap.empty in
  let state_descendants = ref IntMap.empty in
  let transitions = ref IntMap.empty in
  let transition_idx = ref 0 in
  let transition_targets = ref IntMap.empty in
  let datamodel = ref [] in
  (* let is_early = document.Src.Document.binding == `early in *)

  let maybe_map_state idx id =
    match id with
    | None -> ()
    | Some s -> state_map := StringMap.add s idx !state_map
  in

  let rec translate_document () =
    let idx = get_next state_idx in
    let children, on_enter, on_exit, transitions, invocations =
      translate_children document.Src.Document.children idx [] in
    set_mut state_completions idx document.Src.Document.initial;
    set_mut states idx {
      Tgt.State.t=`compound;<commit_msg>fix history completion<commit_after>open Statechart_translator_predicates

module Src = Statechart_t
module Tgt = Statechart_executable
module Bitset = Statechart_bitset

module StringMap = Map.Make(struct
  type t = string
  let compare = compare
end)

module IntMap = Map.Make(struct
  type t = int
  let compare = compare
end)

type child_type = State of int
                | OnEntry of Tgt.expression array
                | OnExit of Tgt.expression array
                | Transition of (unit -> int)
                | Invoke of Tgt.invoke
                | Pass

let get_parent ancestors =
  match ancestors with
  | parent :: _ -> parent
  | _ -> 0

exception UnparsedExpression of string
let unwrap_expr expr =
  match expr with
  | Src.Expr v -> raise (UnparsedExpression v)
  | Src.ExprValue v -> Tgt.string_expr v
  | Src.ExprParsed v -> v
  | Src.ExprUnset -> Tgt.null

let map_to_array a f =
  let size = IntMap.cardinal a in
  Array.init size (fun idx ->
    f (IntMap.find idx a)
  )

let set_mut map idx value =
  map := IntMap.add idx value !map

let append_mut_array arr value =
  arr := Array.append !arr value

let get_next r =
  let idx = !r in
  r := idx + 1;
  idx

let resolve_list map l =
  let targets = List.map (fun k -> StringMap.find k map) l in
  Array.of_list targets

let maybe_find map k =
  try
    Some (IntMap.find k map)
  with
  | _ -> None

let has_history children state_map =
  let n = Array.length children in
  let rec find i =
    if i >= n then false
    else (
      let {Tgt.State.t} = IntMap.find i state_map in
      (t == `history_deep || t == `history_shallow) || find (i + 1)
    )
  in
  find 0

let filter_history_completion history states children =
  children
  |> Bitset.filter (fun i ->
    i != history && (
      let state = IntMap.find i states in
      let t = state.Tgt.State.t in
      t != `history_shallow && t != `history_deep
    )
  )
  |> Bitset.to_idx_array

let find_initial states children =
  match Array.length children with
  | 0 -> [||]
  | n ->
    let rec find i =
      if i >= n then [|Array.get children 0|]
      else (
        let idx = Array.get children i in
        match IntMap.find idx states with
        | {Tgt.State.t=`initial} -> [|idx|]
        | _ -> find (i + 1)
      )
    in
    find 0

let get_completion state initials state_map states descendants =
  match state with
  | {Tgt.State.t=`history_deep; idx; parent} ->
    Array.get descendants parent
    |> filter_history_completion idx states
  | {Tgt.State.t=`history_shallow; idx; parent} ->
    (IntMap.find parent states).Tgt.State.children
    |> Bitset.of_idx_array (IntMap.cardinal states)
    |> filter_history_completion idx states
  | {Tgt.State.t=`parallel; children} ->
    children
  | {Tgt.State.idx=idx} -> (
    match maybe_find initials idx with
    | Some initial when (initial != []) -> resolve_list state_map initial
    | _ -> find_initial states state.Tgt.State.children
  )

let flatten document =
  let states = ref IntMap.empty in
  let state_idx = ref 0 in
  let state_map = ref StringMap.empty in
  let state_completions = ref IntMap.empty in
  let state_descendants = ref IntMap.empty in
  let transitions = ref IntMap.empty in
  let transition_idx = ref 0 in
  let transition_targets = ref IntMap.empty in
  let datamodel = ref [] in
  (* let is_early = document.Src.Document.binding == `early in *)

  let maybe_map_state idx id =
    match id with
    | None -> ()
    | Some s -> state_map := StringMap.add s idx !state_map
  in

  let rec translate_document () =
    let idx = get_next state_idx in
    let children, on_enter, on_exit, transitions, invocations =
      translate_children document.Src.Document.children idx [] in
    set_mut state_completions idx document.Src.Document.initial;
    set_mut states idx {
      Tgt.State.t=`compound;
{'fraction_non_alphanumeric': 0.09141689373297003, 'fraction_numerical': 0.0027247956403269754, 'mean_word_length': 2.6686656671664166, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 57, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4916942', 'n_tokens_mistral': 2436, 'n_tokens_neox': 2366, 'n_words': 783}