(*
 * Copyright (c) 1997-1999 Massachusetts Institute of Technology
 *
 * 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 *
 *)

(* $Id: ast.ml,v 1.7 1999/02/19 17:22:08 athena Exp $ *)

(* Here, we define a representation for a subset of C's abstract
   syntax tree (AST) and provide functions for manipulating it,
   unparsing it, and extracting information. *)

let cvsid = "$Id: ast.ml,v 1.7 1999/02/19 17:22:08 athena Exp $"

(***********************************
 * Program structure 
 ***********************************)
type c_decl = Decl of string * expr
type c_ast =
    Asch of Asched.annotated_schedule
  | Comment of string
  | For of c_ast * c_ast * c_ast * c_ast
  | Block of (c_decl list) * (c_ast list)
  | Binop of string * expr * expr
  | Expr_assign of expr * expr
  | Stmt_assign of expr * expr
  | Comma of c_ast * c_ast


type c_fcn = Fcn of string * string * (c_decl list) * c_ast

let unparse_decl = function
    Decl (a, b) -> a ^ " " ^ unparse_expr b ^ ";\n"


let id = Comment ("Generated by " ^ cvsid)

let foldr_string_concat l = fold_right (^) l ""

let rec unparse_ast = function
    Asch a -> "{\n" ^
      (unparse_annotated a) ^
      "}\n"
  | Comment s -> "  /* " ^ s ^ " */\n"
  | For (a, b, c, d) ->
      "for (" ^
      unparse_ast a ^ "; " ^ unparse_ast b ^ "; " ^ unparse_ast c
      ^ ")" ^ unparse_ast d
  | Block (d, s) ->
      if (s == []) then ""
      else 
	"{\n"                                      ^ 
        foldr_string_concat (map unparse_decl d)   ^ 
        foldr_string_concat (map unparse_ast s)    ^
        "}\n"      
  | Binop (op, a, b) -> (unparse_expr a) ^ op ^ (unparse_expr b)
  | Expr_assign (a, b) -> (unparse_expr a) ^ " = " ^ (unparse_expr b)
  | Stmt_assign (a, b) -> (unparse_expr a) ^ " = " ^ (unparse_expr b) ^ ";\n"
  | Comma (a, b) -> (unparse_ast a) ^ ", " ^ (unparse_ast b)


let unparse_function = function
    Fcn (typ, name, args, body) ->
      let rec unparse_args = function
	  [Decl (a, b)] -> a ^ " " ^ unparse_expr b 
	| (Decl (a, b)) :: s -> a ^ " " ^ unparse_expr b  ^ ", "
	    ^  unparse_args s
	| [] -> ""
      in 
      (typ ^ " " ^ name ^ "(" ^ unparse_args args ^ ")\n" ^
       unparse_ast body)

		

(***************** Extracting Info from ASTs ***************)


(*
 * traverse a a function and return a list of all expressions,
 * in the execution order
 *)
let rec fcn_to_expr_list =
  let rec acode_to_expr_list = function
      AInstr (Assign (_, x)) -> [x]
    | ASeq (a, b) -> 
	(asched_to_expr_list a) @ (asched_to_expr_list b)
    | _ -> []
  and asched_to_expr_list (Annotate (_, _, _, _, code)) =
    acode_to_expr_list code
  and ast_to_expr_list = function
      Asch a -> asched_to_expr_list a
    | Block (_, a) -> flatten (map ast_to_expr_list a)
    | For (_, _, _, body) ->  ast_to_expr_list body
    | _ -> []
	  
  in fun (Fcn (_, _, _, body)) -> ast_to_expr_list body 

      

(***************** Extracting Constants ***************)

(* add a new key & value to a list of (key,value) pairs, where
   the keys are floats and each key is unique up to almost_equal *)

let add_float_key_value list_so_far (k, v) = 
  if exists (fun (k2, v2) -> almost_equal k k2) list_so_far then
    list_so_far
  else
    (k, v) :: list_so_far


(* find all constants in a given expression *)
let rec expr_to_constants = function
  | Real (a, e) -> [(a, e)]
  | Plus a -> flatten (map expr_to_constants a)
  | Times (a, b) -> (expr_to_constants a) @ (expr_to_constants b)
  | Uminus a -> expr_to_constants a
  | FunctionCall (_, a) -> expr_to_constants a
  | _ -> []


let extract_constants f =
  let constlist = flatten (map expr_to_constants (fcn_to_expr_list f))
  in let unique_constants = fold_left add_float_key_value [] constlist
  in let unparsed_constants = foldr_string_concat
      (map 
	 (function (a, e) -> 
	   (konst_of_float a) ^ " = " ^ 
	   (string_of_float a) ^
	   " = \"" ^ (unparse_expr e) ^ "\"\n")
	 unique_constants)
  in  
  "/* List of constants required by this function: \n\n" ^
  unparsed_constants ^
  "\n*/\n"


(***************** Extracting Operation Counts ***************)

let count_stack_vars =
  let rec count_acode = function
    | ASeq (a, b) -> max (count_asched a) (count_asched b)
    | _ -> 0
  and count_asched (Annotate (_, _, decl, _, code)) =
    (length decl) + (count_acode code)
  and count_ast = function
    | Asch a -> count_asched a
    | Block (d, a) -> (length d) + (max_list (map count_ast a))
    | For (_, _, _, body) -> count_ast body
    | _ -> 0
  in function (Fcn (_, _, _, body)) -> count_ast body


let count_memory_acc f =
  let rec count_var_desc = function
    | Array _ -> 1
    | Call (s, v) -> count_var_desc v
    | _ -> 0
  and count_var (v, _) = count_var_desc v
  and count_acode = function
    | AInstr (Assign (v, _)) -> count_var v
    | ASeq (a, b) -> (count_asched a) + (count_asched b)
    | _ -> 0
  and count_asched = function
      Annotate (_, _, _, _, code) -> count_acode code
  and count_ast = function
    | Asch a -> count_asched a
    | Block (_, a) -> (sum_list (map count_ast a))
    | Comma (a, b) -> (count_ast a) + (count_ast b)
    | For (_, _, _, body) -> count_ast body
    | _ -> 0
  and count_acc_expr_func acc = function
    | Var v -> acc + (count_var v)
    | Plus a -> fold_left count_acc_expr_func acc a
    | Times (a, b) -> fold_left count_acc_expr_func acc [a; b]
    | Uminus a -> count_acc_expr_func acc a
    | FunctionCall (_, a) -> count_acc_expr_func acc a
    | _ -> acc
  in let (Fcn (typ, name, args, body)) = f
  in (count_ast body) + 
    fold_left count_acc_expr_func 0 (fcn_to_expr_list f)


let rec count_flops_expr_func (adds, mults) = function
  | Plus [] -> (adds, mults)
  | Plus a -> 
      let (newadds,newmults) = 
	fold_left count_flops_expr_func (adds, mults) a
      in (newadds + (length a) - 1, newmults)
  | Times (a,b) -> 
      let (newadds, newmults) = 
	fold_left count_flops_expr_func (adds, mults) [a; b]
      in (newadds, newmults + 1)
  | Uminus a -> count_flops_expr_func (adds, mults) a
  | FunctionCall (f, a) -> count_flops_expr_func (adds, mults) a
  | _ -> (adds, mults)


let count_flops f = 
    fold_left count_flops_expr_func (0, 0) (fcn_to_expr_list f)

