<commit_before>
(* This file is free software, part of Logtk. See file "license" for more details. *)

(** {1 Feature Vector indexing} *)

(** Feature Vector indexing (see Schulz 2004) for efficient forward
    and backward subsumption *)

module T = Term
module Fmt = CCFormat

type labels = Index_intf.labels

type feature =
  | N of int
  | S of ID.Set.t
  | M of int ID.Map.t
  | L of labels

type feature_vector = feature IArray.t
(** a vector of feature *)

let mk_n i = N i
let mk_s s = S s
let mk_m m = M m
let mk_l s = L s

module Feature : sig
  type t = feature
  include Interfaces.ORD with type t := t
  include Interfaces.EQ with type t := t
  include Interfaces.PRINT with type t := t
  val add : t -> t -> t
  val max : t -> t -> t
  val leq : t -> t -> bool
end = struct
  type t = feature

  let pp out (f:t): unit = match f with
    | N i -> Fmt.int out i
    | S s -> Fmt.fprintf out "(@[set@ %a@])" (Fmt.seq ID.pp) (ID.Set.to_seq s)
    | M m ->
      Fmt.fprintf out "(@[mset@ %a@])"
        Fmt.(seq (pair ~sep:silent ID.pp int)) (ID.Map.to_seq m)
    | L l ->
      Fmt.fprintf out "(@[labels@ %a@])" Fmt.(seq int) (Util.Int_set.to_seq l)

  let to_string = Fmt.to_string pp

  let compare f1 f2 : int =
    let to_int = function N _ -> 0 | S _ -> 1 | M _ -> 2 | L _ -> 3 in
    match f1, f2 with
      | N i1, N i2 -> CCInt.compare i1 i2
      | S s1, S s2 -> ID.Set.compare s1 s2
      | M m1, M m2 -> ID.Map.compare CCInt.compare m1 m2
      | L m1, L m2 -> Util.Int_set.compare m1 m2
      | N _, _
      | S _, _
      | M _, _
      | L _, _
        -> CCInt.compare (to_int f1)(to_int f2)

  let equal f1 f2 = compare f1 f2 = 0

  (* combination of features *)
  let combine_ ~default ~op f1 f2: feature = match f1, f2 with
    | N i1, N i2 -> mk_n (i1+i2)
    | S s1, S s2 -> mk_s (ID.Set.union s1 s2)
    | M m1, M m2 ->
      ID.Map.merge
        (fun _ o1 o2 ->
           Some (op (CCOpt.get_or ~default o1) (CCOpt.get_or ~default o2)))
        m1 m2
      |> mk_m
    | L s1, L s2 -> Util.Int_set.union s1 s2 |> mk_l
    | N _, _
    | S _, _
    | M _, _
    | L _, _
      -> assert false

  let add = combine_ ~default:0 ~op:(+)
  let max = combine_ ~default:0 ~op:max

  let leq (f1:t) (f2:t): bool = match f1, f2 with
    | N i1, N i2 -> i1 <= i2
    | S s1, S s2 -> ID.Set.subset s1 s2
    | M m1, M m2 ->
      ID.Map.for_all
        (fun k i1 -> i1 <= ID.Map.get_or ~default:~-1 k m2)
        m1
    | L s1, L s2 -> Util.Int_set.subset s1 s2
    | N _, _
    | S _, _
    | M _, _
    | L _, _
      -> assert false
end

(** {2 Features} *)

module Make(C: Index_intf.CLAUSE) = struct
  module C = C

  module Feature_fun = struct
    type t = {
      name : string;
      f : Index_intf.lits -> Index_intf.labels -> feature;
    } (** a function that computes a given feature on clauses *)

    let compute f c = f.f (C.to_lits c) (C.labels c)

    let name f = f.name

    let pp out f = Fmt.string out f.name
    let to_string = Fmt.to_string pp

    let make name f = {name; f}

    let size_plus =
      make "size+"
        (fun lits _ ->
           Iter.filter SLiteral.is_pos lits |> Iter.length |> mk_n)

    let size_minus =
      make "size-"
        (fun lits _ ->
           Iter.filter SLiteral.is_neg lits |> Iter.length |> mk_n)

    let weight_lit lit =
      let rec t_weight t = 
        let pref,t = T.open_fun t in
        match T.view t with
        | T.Const _ | T.Var _ | T.DB _ -> 1
        | T.Fun(_,t) -> assert false
        | T.AppBuiltin(_,l)
        | T.App(_,l) ->
          if Term.is_app_var t then List.length pref + 1
          else (List.fold_left (fun acc s -> acc + t_weight s) (List.length pref+1) l) in

      SLiteral.to_seq lit |> Iter.map t_weight |> Iter.fold (+) 0

    let weight_ name filter =
      make name
        (fun lits _ ->
           Iter.filter filter lits
           |> Iter.map weight_lit
           |> Iter.fold (+) 0
           |> mk_n)

    let weight_plus = weight_ "weight+" SLiteral.is_pos
    let weight_minus = weight_ "weight-" SLiteral.is_neg

    let labels = make "labels" (fun _ labels -> mk_l labels)

    let not_app_var t = not (T.is_app_var t)

    (* sequence of symbols of clause, of given sign *)
    let symbols_ filter lits : ID.t Iter.t =
      lits
      |> Iter.filter filter
      |> Iter.flat_map SLiteral.to_seq
      |> Iter.flat_map (T.Seq.symbols ~filter_term:not_app_var)

    let set_sym_ filter lits _ =
      symbols_ filter lits
      |> ID.Set.of_seq
      |> mk_s

    let set_sym_plus =
      make "set_symb+" (set_sym_ SLiteral.is_pos)

    let set_sym_minus =
      make "set_symb-" (set_sym_ SLiteral.is_neg)

    let multiset_sym_ filter lits _ =
      symbols_ filter lits
      |> Iter.fold
        (fun m id ->
           let n = ID.Map.get_or ~default:0 id m in
           ID.Map.add id (n+1) m)
        ID.Map.empty
      |> mk_m

    let multiset_sym_plus =
      make "mset_symb+" (multiset_sym_ SLiteral.is_pos)

    let multiset_sym_minus =
      make "mset_symb-" (multiset_sym_ SLiteral.is_neg)

    (* sequence of symbols of clause with their depth *)
    let symbols_depth_ filter lits : (ID.t * int) Iter.t =
      lits
      |> Iter.filter filter
      |> Iter.flat_map SLiteral.to_seq
      |> Iter.flat_map (T.Seq.subterms_depth ~filter_term:not_app_var)
      |> Iter.filter_map
        (fun (t,d) -> match T.view t with
           | T.Const id -> Some (id,d)
           | _ -> None)

    let depth_sym_ filter lits _ =
      symbols_depth_ filter lits
      |> Iter.fold
        (fun m (id, d) ->
           let d' = ID.Map.get_or ~default:0 id m in
           ID.Map.add id (max d d') m)
        ID.Map.empty
      |> mk_m

    let depth_sym_plus =
      make "depth_sym+" (depth_sym_ SLiteral.is_pos)

    let depth_sym_minus =
      make "depth_sym-" (depth_sym_ SLiteral.is_neg)
  end

  type feature_funs = Feature_fun.t IArray.t

  let compute_fv funs lits labels : feature_vector =
    IArray.map (fun feat -> feat.Feature_fun.f lits labels) funs

  let compute_fv_c funs c : feature_vector = compute_fv funs (C.to_lits c) (C.labels c)

  (** {2 Feature Trie} *)<commit_msg>more fixes to FV tree indexing<commit_after>
(* This file is free software, part of Logtk. See file "license" for more details. *)

(** {1 Feature Vector indexing} *)

(** Feature Vector indexing (see Schulz 2004) for efficient forward
    and backward subsumption *)

module T = Term
module Fmt = CCFormat

type labels = Index_intf.labels

type feature =
  | N of int
  | S of ID.Set.t
  | M of int ID.Map.t
  | L of labels

type feature_vector = feature IArray.t
(** a vector of feature *)

let mk_n i = N i
let mk_s s = S s
let mk_m m = M m
let mk_l s = L s

module Feature : sig
  type t = feature
  include Interfaces.ORD with type t := t
  include Interfaces.EQ with type t := t
  include Interfaces.PRINT with type t := t
  val add : t -> t -> t
  val max : t -> t -> t
  val leq : t -> t -> bool
end = struct
  type t = feature

  let pp out (f:t): unit = match f with
    | N i -> Fmt.int out i
    | S s -> Fmt.fprintf out "(@[set@ %a@])" (Fmt.seq ID.pp) (ID.Set.to_seq s)
    | M m ->
      Fmt.fprintf out "(@[mset@ %a@])"
        Fmt.(seq (pair ~sep:silent ID.pp int)) (ID.Map.to_seq m)
    | L l ->
      Fmt.fprintf out "(@[labels@ %a@])" Fmt.(seq int) (Util.Int_set.to_seq l)

  let to_string = Fmt.to_string pp

  let compare f1 f2 : int =
    let to_int = function N _ -> 0 | S _ -> 1 | M _ -> 2 | L _ -> 3 in
    match f1, f2 with
      | N i1, N i2 -> CCInt.compare i1 i2
      | S s1, S s2 -> ID.Set.compare s1 s2
      | M m1, M m2 -> ID.Map.compare CCInt.compare m1 m2
      | L m1, L m2 -> Util.Int_set.compare m1 m2
      | N _, _
      | S _, _
      | M _, _
      | L _, _
        -> CCInt.compare (to_int f1)(to_int f2)

  let equal f1 f2 = compare f1 f2 = 0

  (* combination of features *)
  let combine_ ~default ~op f1 f2: feature = match f1, f2 with
    | N i1, N i2 -> mk_n (i1+i2)
    | S s1, S s2 -> mk_s (ID.Set.union s1 s2)
    | M m1, M m2 ->
      ID.Map.merge
        (fun _ o1 o2 ->
           Some (op (CCOpt.get_or ~default o1) (CCOpt.get_or ~default o2)))
        m1 m2
      |> mk_m
    | L s1, L s2 -> Util.Int_set.union s1 s2 |> mk_l
    | N _, _
    | S _, _
    | M _, _
    | L _, _
      -> assert false

  let add = combine_ ~default:0 ~op:(+)
  let max = combine_ ~default:0 ~op:max

  let leq (f1:t) (f2:t): bool = match f1, f2 with
    | N i1, N i2 -> i1 <= i2
    | S s1, S s2 -> ID.Set.subset s1 s2
    | M m1, M m2 ->
      ID.Map.for_all
        (fun k i1 -> i1 <= ID.Map.get_or ~default:~-1 k m2)
        m1
    | L s1, L s2 -> Util.Int_set.subset s1 s2
    | N _, _
    | S _, _
    | M _, _
    | L _, _
      -> assert false
end

(** {2 Features} *)

module Make(C: Index_intf.CLAUSE) = struct
  module C = C

  module Feature_fun = struct
    type t = {
      name : string;
      f : Index_intf.lits -> Index_intf.labels -> feature;
    } (** a function that computes a given feature on clauses *)

    let compute f c = f.f (C.to_lits c) (C.labels c)

    let name f = f.name

    let pp out f = Fmt.string out f.name
    let to_string = Fmt.to_string pp

    let make name f = {name; f}

    let size_plus =
      make "size+"
        (fun lits _ ->
           Iter.filter SLiteral.is_pos lits |> Iter.length |> mk_n)

    let size_minus =
      make "size-"
        (fun lits _ ->
           Iter.filter SLiteral.is_neg lits |> Iter.length |> mk_n)

    let weight_lit lit =
      let rec t_weight t = 
        let pref,t = T.open_fun t in
        match T.view t with
        | T.Const _ | T.Var _ | T.DB _ -> List.length pref + 1
        | T.Fun(_,t) -> assert false
        | T.AppBuiltin(_,l)
        | T.App(_,l) ->
          if Term.is_app_var t then List.length pref + 1
          else (List.fold_left (fun acc s -> acc + t_weight s) (List.length pref + 1) l) in

      SLiteral.to_seq lit |> Iter.map t_weight |> Iter.fold (+) 0

    let weight_ name filter =
      make name
        (fun lits _ ->
           Iter.filter filter lits
           |> Iter.map weight_lit
           |> Iter.fold (+) 0
           |> mk_n)

    let weight_plus = weight_ "weight+" SLiteral.is_pos
    let weight_minus = weight_ "weight-" SLiteral.is_neg

    let labels = make "labels" (fun _ labels -> mk_l labels)

    let not_app_var t = not (T.is_app_var t)

    (* sequence of symbols of clause, of given sign *)
    let symbols_ filter lits : ID.t Iter.t =
      lits
      |> Iter.filter filter
      |> Iter.flat_map SLiteral.to_seq
      |> Iter.flat_map (T.Seq.symbols ~filter_term:not_app_var)

    let set_sym_ filter lits _ =
      symbols_ filter lits
      |> ID.Set.of_seq
      |> mk_s

    let set_sym_plus =
      make "set_symb+" (set_sym_ SLiteral.is_pos)

    let set_sym_minus =
      make "set_symb-" (set_sym_ SLiteral.is_neg)

    let multiset_sym_ filter lits _ =
      symbols_ filter lits
      |> Iter.fold
        (fun m id ->
           let n = ID.Map.get_or ~default:0 id m in
           ID.Map.add id (n+1) m)
        ID.Map.empty
      |> mk_m

    let multiset_sym_plus =
      make "mset_symb+" (multiset_sym_ SLiteral.is_pos)

    let multiset_sym_minus =
      make "mset_symb-" (multiset_sym_ SLiteral.is_neg)

    (* sequence of symbols of clause with their depth *)
    let symbols_depth_ filter lits : (ID.t * int) Iter.t =
      (* ignores app vars and does not count opening the functions *)
      let subterms_depth t k =
      let rec recurse depth t =
        if not (T.is_app_var t) then  (
          k (t, depth);
          match T.view t with
            | Const _
            | DB _
            | Var _ -> ()
            | Fun (_,u) -> recurse depth u (* not increasing the depth when we go under a fun 
                                                -- no need to eta expand *)
            | AppBuiltin (_, l) -> List.iter (recurse (depth+1)) l
            | App (hd, l) ->
              recurse depth hd;
              List.iter (recurse (depth+1)) l)
      in
      recurse 0 t in
    
      lits
      |> Iter.filter filter
      |> Iter.flat_map SLiteral.to_seq
      |> Iter.flat_map (subterms_depth)
      |> Iter.filter_map
        (fun (t,d) -> match T.view t with
           | T.Const id -> Some (id,d)
           | _ -> None)

    let depth_sym_ filter lits _ =
      symbols_depth_ filter lits
      |> Iter.fold
        (fun m (id, d) ->
           let d' = ID.Map.get_or ~default:0 id m in
           ID.Map.add id (max d d') m)
        ID.Map.empty
      |> mk_m

    let depth_sym_plus =
      make "depth_sym+" (depth_sym_ SLiteral.is_pos)

    let depth_sym_minus =
      make "depth_sym-" (depth_sym_ SLiteral.is_neg)
  end

  type feature_funs = Feature_fun.t IArray.t

  let compute_fv funs lits labels : feature_vector =
    IArray.map (fun feat -> feat.Feature_fun.f lits labels) funs

  let compute_fv_c funs c : feature_vector = compute_fv funs (C.to_lits c) (C.labels c)

  (** {2 Feature Trie} *)
{'fraction_non_alphanumeric': 0.11540832049306625, 'fraction_numerical': 0.014869029275808937, 'mean_word_length': 1.9014304872597227, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 149, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3175883', 'n_tokens_mistral': 4951, 'n_tokens_neox': 4792, 'n_words': 1522}