(***********************************************************************)
(*                                                                     *)
(*                               TAToo                                 *)
(*                                                                     *)
(*                     Kim Nguyen, LRI UMR8623                         *)
(*                   Université Paris-Sud & CNRS                       *)
(*                                                                     *)
(*  Copyright 2010-2013 Université Paris-Sud and Centre National de la *)
(*  Recherche Scientifique. All rights reserved.  This file is         *)
(*  distributed under the terms of the GNU Lesser General Public       *)
(*  License, with the special exception on linking described in file   *)
(*  ../LICENSE.                                                        *)
(*                                                                     *)
(***********************************************************************)

open Ast


let ( => ) a b = (a, b)
let ( ++ ) a b = Ata.Formula.or_ a b
let ( %% ) a b = Ata.Formula.and_ a b
let ( @: ) a b = StateSet.add a b

module F = Ata.Formula


let node_set = QNameSet.remove QName.document QNameSet.any
let star_set = QNameSet.diff QNameSet.any (
  List.fold_right (QNameSet.add)
    [ QName.document; QName.text; QName.comment ]
    QNameSet.empty)
let root_set = QNameSet.singleton QName.document

(* [compile_axis_test axis test q phi trans states] Takes an xpath
   [axis] and node [test], a formula [phi], a list of [trans]itions
   and a set of [states] and returns a formula [phi'], a new set of
   transitions, and a new set of states such that [phi'] holds iff
   there exists a node reachable through [axis]::[test] where [phi]
   holds.
*)

let compile_axis_test axis (test,kind) phi trans states =
  let q = State.next () in
  let phi = match kind with
      Tree.NodeKind.Node -> phi
    | _ -> phi %% F.is kind
  in
  let phi', trans', states' =
    match axis with
    | Self ->
        (F.stay q,
         (q, [  test => phi ]) :: trans,
         states)

    | Child ->
        (F.first_child q,
         (q, [ test => phi;
               QNameSet.any => F.next_sibling q ]) :: trans,
         states)

    | Descendant false ->
        (F.first_child q,
         (q, [ test => phi;
               QNameSet.any => F.first_child q ++ F.next_sibling q;
             ]) :: trans,
         states)
    | Descendant true ->
        let q' = State.next () in
        (F.stay q ++ F.first_child q',
         (q', [ QNameSet.any => F.stay q ++ F.first_child q' ++ F.next_sibling q';
              ])::
           (q, [ test => phi]):: trans,
         states)

    | Parent ->
        let q' = State.next () in
        let move = F.parent q ++ F.previous_sibling q' in
        (move,
         (q, [ test => phi ])
         :: (q', [ QNameSet.any => move ]) :: trans,
         (q' @: states))

    | Ancestor self ->
      let q' = State.next () in
      let move = F.parent q' ++ F.previous_sibling q' in
      (if self then F.stay q ++ F.stay q' else F.stay q'),
      (q', [ QNameSet.any => move ++ F.parent q])
      :: (q, [ test => phi ]) :: trans,
      (q' @: states)

    | FollowingSibling | PrecedingSibling ->
        let move =
          if axis = PrecedingSibling then
            F.previous_sibling q
          else F.next_sibling q
        in
        move,
        (q, [ test => phi;
              QNameSet.any => move ]) :: trans,
        states

    | Attribute ->
        (F.first_child q,
         (q, [ test => phi;
               QNameSet.any => F.next_sibling q]) :: trans,
         states)
    | _ -> assert false

  in
  phi', trans', q @: states'

let rec compile_expr e trans states =
  match e with
  | Binop (e1, (And|Or as op), e2) ->
      let phi1, trans1, states1 = compile_expr e1 trans states in
      let phi2, trans2, states2 = compile_expr e2 trans1 states1 in
      (if op = Or then phi1 ++ phi2 else phi1 %% phi2),
      trans2,
      states2
  | Fun_call (f, [ e0 ]) when (QName.to_string f) = "not" ->
      let phi, trans0, states0 = compile_expr e0 trans states in
      (F.not_ phi),
      trans0,
      states0
  | Path p -> compile_path p trans states

  | _ -> assert false
and compile_path paths trans states =
  List.fold_left (fun (aphi, atrans, astates) p ->
    let phi, ntrans, nstates = compile_single_path p atrans astates in
    (F.or_ phi aphi),
    ntrans,
    nstates) (F.false_,trans,states) paths

and compile_single_path p trans states =
  let steps =
    match p with
    | Absolute steps ->
        (Ancestor false, (QNameSet.singleton QName.document,
                          Tree.NodeKind.Node), [])
        :: steps
    | Relative steps -> steps
  in
  compile_step_list steps trans states

and compile_step_list l trans states =
  match l with
  | [] -> F.true_, trans, states
  | (axis, test, elist) :: ll ->
      let phi0, trans0, states0 = compile_step_list ll trans states in
      let phi1, trans1, states1 =
        compile_axis_test axis test phi0 trans0 states0
      in
      List.fold_left (fun (aphi, atrans, astates) e ->
        let ephi, etrans, estates = compile_expr e atrans astates in
        aphi %% ephi, etrans, estates) (phi1, trans1, states1) elist

(**
   Compile the top-level XPath query in reverse (going downward
   to the last top-level state):
   /a0::t0[p0]/../an-1::tn-1[pn-1]/an::tn[pn] becomes:
   self::node()[ pn and
   self::tn[pn]/inv(an)::(tn-1)[pn-1]/.../inv(a1)::t0[p0]/inv(a0)::document()]

   /child::a/attribute::b
   self::@b/parent::a/parent::doc()
*)

let compile_top_level_step_list l trans states =
  let rec loop l trans states phi_above =
    match l with
    | [] -> assert false
    | (axis, (test,kind), elist) :: ll ->
        let phi0, trans0, states0 =
          compile_axis_test (invert_axis axis)
            (QNameSet.any, Tree.NodeKind.Node)
            phi_above trans states
        in
        (* Only select attribute nodes if the previous axis
           is attribute *)
        let phi0 =
          if axis != Attribute && kind == Tree.NodeKind.Node then
            phi0 %% (F.not_ F.is_attribute)
          else phi0
        in
        match ll with
          [] ->
          let phi1, trans1, states1 =
            List.fold_left (fun (aphi, atrans, astates) e ->
                let ephi, etrans, estates = compile_expr e atrans astates in
                aphi %% ephi, etrans, estates) (phi0, trans0, states0) elist
          in
          let _, trans2, states2 =
            compile_axis_test Self (test,kind) phi1 trans1 states1
          in
          let marking_state =
            StateSet.choose (StateSet.diff states2 states1)
          in
          marking_state, trans2, states2
        | _ ->
          let phi1, trans1, states1 =
            compile_axis_test Self (test,kind) phi0 trans0 states0
          in
          let phi2, trans2, states2 =
            List.fold_left (fun (aphi, atrans, astates) e ->
                let ephi, etrans, estates = compile_expr e atrans astates in
                aphi %% ephi, etrans, estates) (phi1, trans1, states1) elist
          in
          loop ll trans2 states2  phi2
  in
  let starting = State.next () in
  let phi0, trans0, states0 =
    compile_axis_test
      Self
      (QNameSet.any, Tree.NodeKind.Node)
      (F.stay starting)
      trans
      states
  in
  let mstates, trans, states = loop l trans0 states0 phi0 in
  starting, mstates, trans, states
;;

let path p =
  let sstates, mstates, trans, states =
    List.fold_left (fun (ass, ams, atrs, asts) p ->
      let ss, ms, natrs, nasts =
        match p with
        | Absolute l | Relative l -> compile_top_level_step_list l atrs asts
      in
      (StateSet.add ss ass),
      (StateSet.add ms ams),
      natrs,
      nasts) (StateSet.empty, StateSet.empty, [], StateSet.empty) p
  in
  let builder = Ata.Builder.make () in
  (** ensure that we have a single selecting state at the end *)
  let phi_sel = StateSet.fold (fun q acc -> F.or_ (F.stay q) acc) mstates F.false_ in
  let q_sel = State.next () in
  let states = StateSet.add q_sel states in
  let mstates = StateSet.singleton q_sel in
  let trans = (q_sel, [QNameSet.any, phi_sel]) :: trans in
  StateSet.iter
    (Ata.Builder.add_state builder ~starting:true) sstates;
  StateSet.iter
    (Ata.Builder.add_state builder ~selecting:true) mstates;
  StateSet.iter
    (Ata.Builder.add_state builder) states;
  List.iter (fun (q, l) ->
    List.iter (fun (lab, phi) ->
      Ata.Builder.add_trans builder q lab phi
    ) l) trans;
  Ata.Builder.finalize builder