(* Memory info command for virtual domains.
   (C) Copyright 2008 Richard W.M. Jones, Red Hat Inc.
   http://libvirt.org/

   This program is free software; you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation; either version 2 of the License, or
   (at your option) any later version.

   This program is distributed in the hope that it will be useful,
   but WITHOUT ANY WARRANTY; without even the implied warranty of
   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
   GNU General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program; if not, write to the Free Software
   Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *)

open Camlp4.PreCast
open Syntax
(*open Ast*)

open ExtList
open ExtString
open Printf

module PP = Pahole_parser
module SC = Struct_classify

let rec uniq ?(cmp = Pervasives.compare) = function
    [] -> []
  | [x] -> [x]
  | x :: y :: xs when cmp x y = 0 ->
      uniq (x :: xs)
  | x :: y :: xs ->
      x :: uniq (y :: xs)

let sort_uniq ?cmp xs =
  let xs = List.sort ?cmp xs in
  let xs = uniq ?cmp xs in
  xs

(* We don't care about locations when generating code, so it's
 * useful to just have a single global _loc.
 *)
let _loc = Loc.ghost

(* Some handy camlp4 construction functions which do some
 * things that ought to be easy/obvious but aren't.
 *
 * 'concat_str_items' concatenates a list of str_item together into
 * one big str_item.
 *
 * 'concat_record_fields' concatenates a list of records fields into
 * a record.  The list must have at least one element.
 *
 * 'build_record' builds a record out of record fields.
 * 
 * 'build_tuple_from_exprs' builds an arbitrary length tuple from
 * a list of expressions of length >= 2.
 *
 * Thanks to bluestorm on #ocaml for getting these working.
 *)
let concat_str_items items =
  match items with
  | [] -> <:str_item< >>
  | x :: xs ->
      List.fold_left (fun xs x -> <:str_item< $xs$ $x$ >>) x xs

let concat_sig_items items =
  match items with
  | [] -> <:sig_item< >>
  | x :: xs ->
      List.fold_left (fun xs x -> <:sig_item< $xs$ $x$ >>) x xs

let concat_exprs exprs =
  match exprs with
  | [] -> assert false
  | x :: xs ->
      List.fold_left (fun xs x -> <:expr< $xs$ ; $x$ >>) x xs

let concat_record_fields fields =
  match fields with
    | [] -> assert false
    | f :: fs ->
	List.fold_left (fun fs f -> <:ctyp< $fs$ ; $f$ >>) f fs

let concat_record_bindings rbs =
  match rbs with
    | [] -> assert false
    | rb :: rbs ->
	List.fold_left (fun rbs rb -> <:rec_binding< $rbs$ ; $rb$ >>) rb rbs

let build_record rbs =
  Ast.ExRec (_loc, rbs, Ast.ExNil _loc)

let build_tuple_from_exprs exprs =
  match exprs with
  | [] | [_] -> assert false
  | x :: xs ->
      Ast.ExTup (_loc,
		 List.fold_left (fun xs x -> Ast.ExCom (_loc, x, xs)) x xs)

let build_tuple_from_patts patts =
  match patts with
  | [] | [_] -> assert false
  | x :: xs ->
      Ast.PaTup (_loc,
		 List.fold_left (fun xs x -> Ast.PaCom (_loc, x, xs)) x xs)

type code = Ast.str_item * Ast.sig_item

let ocaml_type_of_field_type = function
  | PP.FInteger -> <:ctyp< int64 >>
  | PP.FString _ -> <:ctyp< string >>
  | PP.FStructPointer _ | PP.FVoidPointer
  | PP.FAnonListHeadPointer | PP.FListHeadPointer _ ->
      <:ctyp< Virt_mem_mmap.addr >>

let generate_types xs =
  let strs = List.map (
    fun (struct_name, sflist, cflist) ->
      let sflist = List.map (
	fun { SC.sf_name = sf_name; sf_fields = fields } ->
	  if fields <> [] then (
	    let fields = List.map (
	      fun (name, t) ->
		let t = ocaml_type_of_field_type t in
		<:ctyp< $lid:sf_name^"_"^name$ : $t$ >>
	    ) fields in
	    let fields = concat_record_fields fields in

	    <:str_item<
	      type $lid:sf_name$ = { $fields$ }
            >>
	  ) else
	    <:str_item< type $lid:sf_name$ = unit >>
      ) sflist in
      let sflist = concat_str_items sflist in

      let cflist = List.map (
	fun { SC.cf_name = cf_name; cf_fields = fields } ->
	  if fields <> [] then (
	    let fields = List.map (
	      fun (name, t) ->
		let t = ocaml_type_of_field_type t in
		<:ctyp< $lid:cf_name^"_"^name$ : $t$ >>
	    ) fields in
	    let fields = concat_record_fields fields in

	    <:str_item<
	      type $lid:cf_name$ = { $fields$ }
            >>
	  ) else
	    <:str_item< type $lid:cf_name$ = unit >>
      ) cflist in
      let cflist = concat_str_items cflist in

      <:str_item<
        type ('a, 'b) $lid:struct_name$ = 'a * 'b ;;
	$sflist$
	$cflist$
      >>
  ) xs in

  concat_str_items strs, <:sig_item< >>

let generate_offsets xs =
  (* Only need to generate the offset_of_* functions for fields
   * which are cross-referenced from another field.  Which
   * ones are those?
   *)
  let fields =
    List.concat (
      List.map (
	fun (_, (_, all_fields)) ->
	  List.filter_map (
	    function
	    | (_,
	       PP.FListHeadPointer ((Some (struct_name, field_name)) as f)) ->
		f
	    | _ ->
		None
	  ) all_fields
      ) xs
    ) in

  let fields = sort_uniq fields in

  let strs =
    List.map (
      fun (struct_name, field_name) ->
	let kernels, _ =
	  try List.assoc struct_name xs
	  with Not_found ->
	    failwith (
	      sprintf "generate_offsets: structure %s not found. This is probably a list_head-related bug."
		struct_name
	    ) in
	(* Find the offset of this field in each kernel version. *)
	let offsets =
	  List.filter_map (
	    fun ({ PP.kernel_version = version },
		 { PP.struct_fields = fields }) ->
	      try
		let field =
		  List.find (fun { PP.field_name = name } -> field_name = name)
		    fields in
		let offset = field.PP.field_offset in
		Some (version, offset)
	      with Not_found -> None
	  ) kernels in

	if offsets = [] then
	  failwith (
	    sprintf "generate_offsets: field %s.%s not found in any kernel. This is probably a list_head-related bug."
	      struct_name field_name
	  );

	(* Generate a map of kernel version to offset. *)
	let map = List.fold_left (
	  fun map (version, offset) ->
	    <:expr< StringMap.add $str:version$ $`int:offset$ $map$ >>
	) <:expr< StringMap.empty >> offsets in

	let code =
	  <:str_item<
	    let $lid:"offset_of_"^struct_name^"_"^field_name$ =
	      let map = $map$ in
	      fun kernel_version -> StringMap.find kernel_version map
	  >> in
	code
    ) fields in

  let strs = concat_str_items strs in

  strs, <:sig_item< >>

let generate_parsers xs =
  let strs =
    List.map (
      fun (struct_name, palist) ->
	let palist =
	  List.map (
	    fun { SC.pa_name = pa_name } ->
	      <:str_item<
		let $lid:pa_name$ kernel_version bits = $str:pa_name$
	      >>
	  ) palist in
	concat_str_items palist
    ) xs in

  let strs = concat_str_items strs in

  (* The shared parser functions.
   * 
   * We could include bitmatch statements directly in here, but
   * what happens is that the macros get expanded here, resulting
   * in (even more) unreadable generated code.  So instead just
   * do a textual substitution later by post-processing the
   * generated files.  Not type-safe, but we can't have
   * everything.
   *)
  let subs = Hashtbl.create 13 in
  List.iter (
    fun (struct_name, palist) ->
      List.iter (
	fun ({ SC.pa_name = pa_name;
	       pa_endian = endian; pa_structure = structure;
	       pa_shape_field_struct = sf;
	       pa_content_field_struct = cf }) ->
	  (* Generate the code to match this structure. *)
	  let endian =
	    match endian with
	    | Bitstring.LittleEndian -> "littleendian"
	    | Bitstring.BigEndian -> "bigendian"
	    | _ -> assert false in
	  let patterns =
	    String.concat ";\n      " (
	      List.map (
		function
		| { PP.field_name = field_name;
		    field_type = PP.FInteger;
		    field_offset = offset;
		    field_size = size } ->
		    (* 'zero+' is a hack to force the type to int64. *)
		    sprintf "%s : zero+%d : offset(%d), %s"
		      field_name (size*8) (offset*8) endian

		| { PP.field_name = field_name;
		    field_type = (PP.FStructPointer _
		                  | PP.FVoidPointer
				  | PP.FAnonListHeadPointer
				  | PP.FListHeadPointer _);
		    field_offset = offset;
		    field_size = size } ->
		    sprintf "%s : zero+%d : offset(%d), %s"
		      field_name (size*8) (offset*8) endian

		| { PP.field_name = field_name;
		    field_type = PP.FString width;
		    field_offset = offset;
		    field_size = size } ->
		    sprintf "%s : %d : offset(%d), string"
		      field_name (width*8) (offset*8)
	      ) structure.PP.struct_fields
	    ) in

	  let shape_assignments =
	    List.map (
	      fun (field_name, field_type) ->

		(* Go and look up the field offset in the correct kernel. *)
		let { PP.field_offset = offset } =
		  List.find (fun { PP.field_name = name } -> field_name = name)
		    structure.PP.struct_fields in

		(* Generate assignment code, if necessary we can adjust
		 * the list_head.
		 *)
		match field_type with
		| PP.FListHeadPointer None ->
		    sprintf "%s_%s = (if %s <> 0L then Int64.sub %s %dL else %s)"
		      sf.SC.sf_name field_name
		      field_name
		      field_name offset field_name

		| PP.FListHeadPointer (Some (other_struct_name,
					     other_field_name)) ->
		    (* A reference to a field in another structure.  We don't
		     * know the offset until runtime, so we have to call
		     * offset_of_<struct>_<field> to find it.
		     *)
		    sprintf "%s_%s = (
                      if %s <> 0L then (
                        let offset = offset_of_%s_%s kernel_version in
                        let offset = Int64.of_int offset in
                        Int64.sub %s offset
                      ) else %s
                    )"
		      sf.SC.sf_name field_name field_name
		      other_struct_name other_field_name
		      field_name field_name
		| _ ->
		    sprintf "%s_%s = %s" sf.SC.sf_name field_name field_name
	    ) sf.SC.sf_fields in

	  let shape_assignments =
	    if shape_assignments = [] then "()"
	    else
	      "{ " ^ String.concat ";\n        " shape_assignments ^ " }" in

	  let content_assignments =
	    List.map (
	      fun (field_name, _) ->
		sprintf "%s_%s = %s" cf.SC.cf_name field_name field_name
	    ) cf.SC.cf_fields in

	  let content_assignments =
	    if content_assignments = [] then "()"
	    else
	      "{ " ^ String.concat ";\n        " content_assignments ^ " }" in

	  let code =
	    sprintf "
  bitmatch bits with
  | { %s } ->
      let s =
      %s in
      let c =
      %s in
      (s, c)
  | { _ } ->
      raise (Virt_mem_types.ParseError (%S, %S, match_err))"
	      patterns shape_assignments content_assignments
	      struct_name pa_name in

	  Hashtbl.add subs pa_name code
      ) palist;
  ) xs;

  (strs, <:sig_item< >>), subs

(* Helper functions to store things in a fixed-length tuple very efficiently.
 * Note that the tuple length must be >= 2.
 *)
type tuple = string list

let tuple_create fields : tuple = fields

(* Generates 'let _, _, resultpatt, _ = tupleexpr in body'. *)
let tuple_generate_extract fields field resultpatt tupleexpr body =
  let patts = List.map (
    fun name -> if name = field then resultpatt else <:patt< _ >>
  ) fields in
  let result = build_tuple_from_patts patts in
  <:expr< let $result$ = $tupleexpr$ in $body$ >>

(* Generates '(fieldexpr1, fieldexpr2, ...)'. *)
let tuple_generate_construct fieldexprs =
  build_tuple_from_exprs fieldexprs

type follower_t =
  | Missing of string | Follower of string | KernelVersion of string

let generate_followers xs =
  (* Tuple of follower functions, just a list of struct_names. *)
  let follower_tuple = tuple_create (List.map fst xs) in

  (* A shape-follow function for every structure/shape. *)
  let strs = List.map (
    fun (struct_name, (_, sflist, _, _)) ->
      List.map (
	fun { SC.sf_name = sf_name; sf_fields = fields } ->
	  let body = List.fold_right (
	    fun (name, typ) body ->
	      let follower_name =
		match typ with
		| PP.FListHeadPointer None -> struct_name
		| PP.FListHeadPointer (Some (struct_name, _)) -> struct_name
		| PP.FStructPointer struct_name -> struct_name
		| _ -> assert false in
	      tuple_generate_extract follower_tuple follower_name
	        <:patt< f >> <:expr< followers >>
	        <:expr<
		  let map =
		    f load followers map shape.$lid:sf_name^"_"^name$ in $body$
	        >>
	  ) fields <:expr< map >> in

	  <:str_item<
	    let $lid:sf_name^"_follower"$ load followers map shape =
	      $body$
	  >>
      ) sflist
  ) xs in
  let strs = List.concat strs in

  (* A follower function for every kernel version / structure.  When this
   * function is called starting at some known root, it will load every
   * reachable kernel structure.
   *)
  let strs =
    let common =
      (* Share as much common code as possible to minimize generated
       * code size and benefit i-cache.
       *)
      <:str_item<
	let kv_follower kernel_version struct_name total_size
	    parserfn followerfn
	    load followers map addr =
	  if addr <> 0L && not (AddrMap.mem addr map) then (
	    let map = AddrMap.add addr (struct_name, total_size) map in
	    let bits = load struct_name addr total_size in
	    let shape, _ = parserfn kernel_version bits in
	    followerfn load followers map shape
	  )
	  else map
      >> in

    let fs =
      List.map (
	fun (struct_name, (kernels, _, sfhash, pahash)) ->
	  List.map (
	    fun ({ PP.kernel_version = version; kv_i = kv_i },
		 { PP.struct_total_size = total_size }) ->
	      let { SC.pa_name = pa_name } = Hashtbl.find pahash version in
	      let { SC.sf_name = sf_name } = Hashtbl.find sfhash version in

	      let fname = sprintf "%s_kv%d_follower" struct_name kv_i in

	      <:str_item<
		let $lid:fname$ =
		  kv_follower
		    $str:version$ $str:struct_name$ $`int:total_size$
		    $lid:pa_name$ $lid:sf_name^"_follower"$
	      >>
	  ) kernels
      ) xs in

    let strs = strs @ [ common ] @ List.concat fs in
    strs in

  (* A map from kernel versions to follower functions.
   *
   * For each struct, we have a list of kernel versions which contain
   * that struct.  Some kernels are missing a particular struct, so
   * that is turned into a ParseError exception.
   *)
  let strs =
    let nr_kernels =
      List.fold_left max 0
	(List.map (fun (_, (kernels, _, _, _)) -> List.length kernels) xs) in
    let nr_structs = List.length xs in
    let array = Array.make_matrix nr_kernels (nr_structs+1) (Missing "") in
    List.iteri (
      fun si (struct_name, _) ->
	for i = 0 to nr_kernels - 1 do
	  array.(i).(si+1) <- Missing struct_name
	done
    ) xs;
    List.iteri (
      fun si (struct_name, (kernels, _, _, _)) ->
	List.iter (
	  fun ({ PP.kernel_version = version; kv_i = kv_i }, _) ->
	    array.(kv_i).(0) <- KernelVersion version;
	    array.(kv_i).(si+1) <-
	      Follower (sprintf "%s_kv%d_follower" struct_name kv_i)
	) kernels
    ) xs;

    let array = Array.map (
      fun row ->
	match Array.to_list row with
	| [] | (Missing _|Follower _) :: _ -> assert false
	| KernelVersion kernel_version :: followers -> kernel_version, followers
    ) array in

    let map = List.fold_left (
      fun map (kernel_version, followers) ->
	let followers = List.map (
	  function
	  | Follower fname ->
	      <:expr< $lid:fname$ >>

	  (* no follower for this kernel/struct combination *)
	  | Missing struct_name ->
	      <:expr<
		fun _ _ _ _ ->
		  raise (
		    Virt_mem_types.ParseError (
		      $str:struct_name$, "follower_map", struct_missing_err
		    )
		  )
	      >>
	  | KernelVersion _ -> assert false
	) followers in
	let followers = tuple_generate_construct followers in

	<:expr< StringMap.add $str:kernel_version$ $followers$ $map$ >>
    ) <:expr< StringMap.empty >> (Array.to_list array) in

    let str =
      <:str_item<
	let follower_map = $map$
      >> in
    strs @ [ str ] in

  (* Finally a publicly exposed follower function. *)
  let strs =
    let fs =
      List.map (
	fun (struct_name, (kernels, _, _, _)) ->
	  let fname = sprintf "%s_follower" struct_name in

	  let body =
	    tuple_generate_extract follower_tuple struct_name
	      <:patt< f >> <:expr< followers >>
	      <:expr<
	        f load followers AddrMap.empty addr
	      >> in

	  <:str_item<
	    let $lid:fname$ kernel_version load addr =
	      let followers =
		try StringMap.find kernel_version follower_map
		with Not_found ->
		  unknown_kernel_version kernel_version $str:struct_name$ in
	      $body$
	  >>
      ) xs in

    strs @ fs in

  let sigs =
    List.map (
      fun (struct_name, _) ->
	<:sig_item<
          val $lid:struct_name^"_follower"$ :
	    kernel_version ->
	    (string -> Virt_mem_mmap.addr -> int -> Bitstring.bitstring) ->
	    Virt_mem_mmap.addr ->
	    (string * int) AddrMap.t
	  >>
    ) xs in

  concat_str_items strs, concat_sig_items sigs

let output_interf ~output_file types offsets parsers followers =
  (* Some standard code that appears at the top of the interface file. *)
  let prologue =
    <:sig_item<
      module AddrMap : sig
	type key = Virt_mem_mmap.addr
	type 'a t = 'a Map.Make(Int64).t
	val empty : 'a t
	val is_empty : 'a t -> bool
	val add : key -> 'a -> 'a t -> 'a t
	val find : key -> 'a t -> 'a
	val remove : key -> 'a t -> 'a t
	val mem : key -> 'a t -> bool
	val iter : (key -> 'a -> unit) -> 'a t -> unit
	val map : ('a -> 'b) -> 'a t -> 'b t
	val mapi : (key -> 'a -> 'b) -> 'a t -> 'b t
	val fold : (key -> 'a -> 'b -> 'b) -> 'a t -> 'b -> 'b
	val compare : ('a -> 'a -> int) -> 'a t -> 'a t -> int
	val equal : ('a -> 'a -> bool) -> 'a t -> 'a t -> bool
      end ;;
      type kernel_version = string ;;
    >> in

  let sigs =
    concat_sig_items [ prologue; types; offsets; parsers; followers ] in
  Printers.OCaml.print_interf ~output_file sigs

(* Finally generate the output files. *)
let re_subst = Pcre.regexp "^(.*)\"(\\w+_parser_\\d+)\"(.*)$"

let output_implem ~output_file types offsets parsers parser_subs followers =
  (* Some standard code that appears at the top of the implementation file. *)
  let prologue =
    <:str_item<
      module StringMap = Map.Make (String) ;;
      module AddrMap = Map.Make (Int64) ;;
      type kernel_version = string ;;

      let match_err = "failed to match kernel structure" ;;
      let struct_missing_err = "struct does not exist in this kernel version" ;;

      let unknown_kernel_version version struct_name =
	invalid_arg (Printf.sprintf "%s: unknown kernel version or
struct %s is not supported in this kernel.
Try a newer version of virt-mem, or if the guest is not from a
supported Linux distribution, see this page about adding support:
  http://et.redhat.com/~rjones/virt-mem/faq.html\n"
		       version struct_name) ;;

      let zero = 0 ;;
    >> in

  let strs =
    concat_str_items [ prologue; types; offsets; parsers; followers ] in

  (* Write the new implementation to .ml.new file. *)
  let new_output_file = output_file ^ ".new" in
  Printers.OCaml.print_implem ~output_file:new_output_file strs;

  (* Substitute the parser bodies in the output file. *)
  let ichan = open_in new_output_file in
  let ochan = open_out output_file in

  output_string ochan "\
(* WARNING: This file and the corresponding mli (interface) are
 * automatically generated by the extract/codegen/ program.
 *
 * Any edits you make to this file will be lost.
 *
 * To update this file from the latest kernel database, it is recommended
 * that you do 'make update-kernel-structs'.
 *)\n\n";

  let rec loop () =
    let line = input_line ichan in
    let line =
      if Pcre.pmatch ~rex:re_subst line then (
	let subs = Pcre.exec ~rex:re_subst line in
	let start = Pcre.get_substring subs 1 in
	let template = Pcre.get_substring subs 2 in
	let rest = Pcre.get_substring subs 3 in
	let sub =
	  try Hashtbl.find parser_subs template
	  with Not_found -> assert false in
	start ^ sub ^ rest
      ) else line in
    output_string ochan line; output_char ochan '\n';
    loop ()
  in
  (try loop () with End_of_file -> ());

  close_out ochan;
  close_in ichan;

  Unix.unlink new_output_file