#ifndef SMPFUNCTION_H
#define SMPFUNCTION_H 1

// SMPFunction.h
//
// This header defines the interfaces needed for analyzing functions, performing live variable analysis,
//  putting code into SSA form, etc.

#include <utility>
#include <list>
#include <vector>
#include <map>
#include <set>

#include <cstddef>

#include <pro.h>
#include <ida.hpp>
#include <ua.hpp>

#include "SMPDataFlowAnalysis.h"
#include "SMPInstr.h"
#include "SMPBasicBlock.h"
#include "ProfilerInformation.h"

class SMPProgram;	// forward declaration so we can declare a pointer to an SMPProgram

using namespace std;

// Use IDA info for switch tables to link indirect jumps to successor blocks?
#define SMP_USE_SWITCH_TABLE_INFO 1

struct LocalVar {
	char VarName[MAXSMPVARSTR];
	long offset;
	size_t size;
};

// Entry for each byte address in the stack frame
struct StackFrameEntry {
	struct LocalVar *VarPtr;  // LocalVar that includes this offset
	long offset;  // offset relative to incoming stack pointer
	bool Read;    // was this entry ever read by an instruction?
	bool Written; // was this entry ever written by an instruction?
	bool AddressTaken; // did this entry have its address taken?
	bool ESPRelativeAccess; // ever accessed by ESP+const?
	bool EBPRelativeAccess; // ever accessed by EBP-const? (only if UseFP)
};

enum FuncType { 
	FUNC_UNKNOWN = 0,
	FUNC_SAFE = 1,
	FUNC_UNSAFE = 2
};


// Class encapsulating all that the SMP static analyzer cares to know
//  about a function.
class SMPFunction {
public:
	// Constructors
	SMPFunction(func_t *Info, SMPProgram *pgm);  // Default constructor

	// Get methods
	inline SMPProgram *GetProg(void) const { return Program; };
	inline const char *GetFuncName(void) const { return FuncName; };
	inline long GetTypedDefs(void) const { return TypedDefs; };
	inline long GetUntypedDefs(void) const { return UntypedDefs; };
	inline long GetTypedPhiDefs(void) const { return TypedPhiDefs; };
	inline long GetUntypedPhiDefs(void) const { return UntypedPhiDefs; };
	inline long GetSafeBlocks(void) const { return SafeBlocks; };
	inline long GetUnsafeBlocks(void) const { return UnsafeBlocks; };
	inline set<op_t, LessOp>::iterator GetFirstGlobalName(void) { return GlobalNames.begin(); };
	inline set<op_t, LessOp>::iterator GetLastGlobalName(void) { return GlobalNames.end(); };
	inline size_t NumGlobalNames(void) { return GlobalNames.size(); };
	inline set<op_t, LessOp>::iterator FindGlobalName(op_t SearchOp) { return GlobalNames.find(SearchOp); };
	inline FuncType GetReturnAddressStatus(void) const { return ReturnAddrStatus;}
	inline const ea_t GetStartAddr() const { return FuncInfo.startEA; }; // exposing the start address of the function. Used in RecurseAndMark
	inline const vector<ea_t> GetCallTargets() const { return AllCallTargets; };
	bool GetIsSpeculative() { return IsSpeculative; }
	inline size_t GetNumCallers(void) const { return AllCallSources.size(); };

	// Set methods
	inline void IncTypedPhiDefs(void) { ++TypedPhiDefs; return; };
	inline void IncUntypedPhiDefs(void) { ++UntypedPhiDefs; return; };
	inline void DecTypedPhiDefs(void) { --TypedPhiDefs; return; };
	inline void DecUntypedPhiDefs(void) { --UntypedPhiDefs; return; };
	inline void SetReturnAddressStatus(FuncType funcType) {
		ReturnAddrStatus = funcType;
	}
	inline void SetFuncSafe(bool Status) { SafeFunc = Status; return; };
	inline void SetSpecFuncSafe(bool Status) { SpecSafeFunc = Status; return; };
	inline void SetNeedsFrame(bool Status) { NeedsStackReferent = Status; return; };
	inline void SetSpecNeedsFrame(bool Status) { SpecNeedsStackReferent = Status; return; };
	void SetIsSpeculative(bool IsS) { IsSpeculative = IsS; }
	void AddCallSource(ea_t addr);

	// Query methods
	inline bool HasIndirectCalls(void) const { return IndirectCalls; };
	inline bool HasUnresolvedIndirectCalls(void) const { return UnresolvedIndirectCalls; };
	inline bool HasIndirectJumps(void) const { return IndirectJumps; };
	inline bool HasUnresolvedIndirectJumps(void) const { return UnresolvedIndirectJumps; };
	inline bool HasSharedChunks(void) const { return SharedChunks; };
	inline bool HasGoodRTLs(void) const { return BuiltRTLs; };
	inline bool IsLibFunc(void) const { return LibFunc; };
	inline bool IsLeaf(void) const { return (!IndirectCalls && DirectCallTargets.empty()); };
	inline bool IsSafe(void) const { return SafeFunc; };
	inline bool IsSpecSafe(void) const { return SpecSafeFunc; };
	inline bool IsSafeCallee(void) const { return SafeCallee; };
	inline bool IsSpecSafeCallee(void) const { return SafeCallee; };
	inline bool NeedsStackFrame(void) const { return NeedsStackReferent; };
	inline bool SpecNeedsStackFrame(void) const { return SpecNeedsStackReferent; };
	inline bool WritesAboveReturnAddress(void) const { return WritesAboveRA; };
	inline bool IsGlobalName(op_t RefOp) const { return (GlobalNames.end() != GlobalNames.find(RefOp)); };
	inline bool UsesFramePointer(void) const { return UseFP; };

	// Printing methods
	void Dump(void); // debug dump

	// Analysis methods
	void ResetProcessedBlocks(void); // Set Processed flag to false in all blocks
	void Analyze(void);  // Analyze all instructions in function
	void EmitAnnotations(FILE *AnnotFile);
	void RPONumberBlocks(void);
	void SetLinks(void); // Link basic blocks and map instructions to blocks
	void LiveVariableAnalysis(void);  // Perform Live Variable Analysis across all blocks
	void ComputeSSA(void); // Compute SSA form data structures
	void AliasAnalysis(void); // Find memory writes with possible aliases
	void InferTypes(bool FirstIter); // Determine NUMERIC, POINTER, etc. for all operands
	SMPOperandType InferGlobalDefType(op_t RefOp, int SSANum, SMPBasicBlock *DefBlock, bool CallInst);
		// Can DEF type be inferred from all USEs?
	void ApplyProfilerInformation(ProfilerInformation* pi);
	void AnalyzeMetadataLiveness(void); // Is metadata live or dead for each inst
	bool PropagateGlobalMetadata(op_t UseOp, SMPMetadataType Status, int SSANum);
	void FindRedundantMetadata(void); // Do consecutive DEFs have same type?

private:
	// Data
	SMPProgram* Program;		// pointer to the program I'm part of
	func_t FuncInfo;
	char FuncName[MAXSMPSTR];
	int BlockCount; // number of basic blocks in the function
	bool UseFP;  // Does function use a frame pointer?
	bool StaticFunc; // Is function declared static?
	bool LibFunc; // is function a standard library function?
	bool IndirectCalls; // Does function make indirect calls?
	bool UnresolvedIndirectCalls; // Calls could not all be linked to targets
	bool IndirectJumps; // Does function make indirect jumps?
	bool UnresolvedIndirectJumps; // Jumps could not all be linked to targets
	bool SharedChunks; // Does function share a tail chunk with other functions?
	bool CallsAlloca; // Does function allocate stack space after initial allocation?
	bool AnalyzedSP; // Were stack pointer change points successfully analyzed?
	bool BuiltRTLs;  // Were RTLs built succcessfully for all instructions?
	bool SafeFunc;  // Function needs no bounds checking from mmStrata
	bool SpecSafeFunc;  // Function needs no bounds checking from mmStrata
	bool SafeCallee; // SafeFunc AND Caller of this func does not need a stack referent
	bool SpecSafeCallee; // SafeFunc AND Caller of this func does not need a stack referent
	bool WritesAboveRA; // Function writes to stack above the return address
	bool NeedsStackReferent; // mmStrata will need a stack referent to do bounds checking
	bool SpecNeedsStackReferent; // mmStrata will need a stack referent to do bounds checking
	bool HasIndirectWrites; // Function has at least one indirect memory write
	bool OutgoingArgsComputed; // Were able to compute OutgoingArgsSize
	long TypedDefs;  // How many DEFs in instructions were not UNINIT type after InferTypes()
	long UntypedDefs; // How many DEFs in instructions are UNINIT type after InferTypes()
	long TypedPhiDefs;
	long UntypedPhiDefs;
	long SafeBlocks;  // no unsafe memory writes in block; counter
	long UnsafeBlocks;  // possibly unsafe memory write in block; counter
	size_t Size; // Function size in code bytes
	asize_t LocalVarsSize;  // size of local vars region of stack frame
	ushort CalleeSavedRegsSize; // stack size of callee pushed regs
	int RetAddrSize; // size of return address on stack (4 for most machines)
	asize_t IncomingArgsSize; // size of incoming args on stack
	size_t OutgoingArgsSize;  // portion of LocalVarsSize that is really outgoing args space
	ea_t LocalVarsAllocInstr; // address of instr that allocates stack frame
	ea_t LocalVarsDeallocInstr; // address of epilogue instr that deallocs frame
	adiff_t AllocPointDelta; // IDA sp_delta AFTER stack frame allocation instruction
	sval_t MinStackDelta; // smallest (negative) value that stack pointer reaches, relative
						  // to the value it has at the entry point of the function
	adiff_t LocalVarOffsetLimit; // upper bound on stack-relative offsets
	FuncType  ReturnAddrStatus; // Marked true if the return address is safe from being overwritten
	bool IsSpeculative;	    // Have we started the speculative portion of the analysis for this function.
	list<SMPInstr> Instrs;
	list<SMPBasicBlock> Blocks;
	vector<ea_t> DirectCallTargets; // addresses called directly
	vector<ea_t> IndirectCallTargets; // addresses called by indirect calls
	vector<ea_t> AllCallTargets; // union of direct and indirect
	set<ea_t> AllCallSources; // functions that call this one
	map<ea_t, list<SMPBasicBlock>::iterator> InstBlockMap;
	vector<list<SMPBasicBlock>::iterator> RPOBlocks;
	vector<int> IDom; // Immediate dominators, indexed and valued by block RPO numbers
	vector<pair<int, list<int> > > DomTree; // Dominator tree, as parent # and list of children
	set<op_t, LessOp> GlobalNames;  // operands used in more than one block; needed in SSA
	vector<list<int> > BlocksDefinedIn; // What blocks DEF each GlobalName; index = op # in GlobalNames
	vector<int> SSACounter; // SSA subscript #, indexed by GlobalNames op # 
	vector<list<int> > SSAStack; // SSA stack of most recent SSA number, indexed by global #
	vector<struct LocalVar> LocalVarTable; // offset-sorted list of local vars / outgoing args
	vector<struct StackFrameEntry> StackFrameMap; // memory map of every byte on stack frame
	vector<int> SavedRegLoc; // indexed by reg #; offset from return address of callee-saved reg
	vector<SMPOperandType> ReturnRegTypes; // indexed by reg #; inferred types upon return

	// Methods
	void SetStackFrameInfo(void);
	ea_t FindAllocPoint(asize_t); // Deal with difficult to find stack frame allocations
	bool MDFixFrameInfo(void); // Redefine stack regions for our needs
	bool MDFixUseFP(void);  // Fix IDA errors affecting UseFP
	void BuildLocalVarTable(void); // Determine local variable boundaries on the stack
	void SemiNaiveLocalVarID(void); // Semi-naive algorithm for local var boundaries ID
	void FindOutgoingArgsSize(void); // Find portion of local frame that is outgoing args
	bool WritesToOutgoingArgs(op_t DestOp); // Does DestOp fall within outgoing args area?
	bool WritesAboveLocalFrame(op_t DestOp); // Is DestOp direct stack write to caller's frame?
	bool IndexedWritesAboveLocalFrame(op_t DestOp); // Is DestOp direct stack write to caller's frame?
	bool MDGetStackOffsetAndSize(op_t TempOp, sval_t sp_delta, ea_t &offset, size_t &DataSize,
		bool &FP);  // Find any stack memory access in TempOp, return offset, size, and whether
					// the Frame Pointer was used.
	bool FindAlloca(void); // true if found evidence of alloca() allocations
	void MDFindSavedRegs(void); // Fill in SavedRegLoc[] offsets
	void EmitStackFrameAnnotations(FILE *AnnotFile, list<SMPInstr>::iterator Instr);
	void ComputeIDoms(void); // Compute immediate dominators of all blocks into IDom[]
	int IntersectDoms(int, int) const; // Find Dom intersection (as IDom[] index) for 2 blocks
	void ComputeDomFrontiers(void); // Compute dominance frontiers for all blocks
	void ComputeGlobalNames(void); // Compute the GlobalNames set
	void ComputeBlocksDefinedIn(void); // Compute the BlocksDefinedIn vector
	void InsertPhiFunctions(void); // Insert SSA phi functions at top of each basic block
	void BuildDominatorTree(void); // Build the DomTree structure
	int SSANewNumber(size_t GlobNameIndex); // SSA helper: increment and return SSA number
	void SSARename(int BlockNumber); // SSA main helper: rename throughout block
	void SSARenumber(void); // Renumber SSA subscripts for all names
	bool FindPossibleChainAlias(list<SMPInstr>::iterator CurrInst, op_t DefOp, int SSANum); 
	   // Does the DefOp DEF_USE chain starting at CurrInst contain an indirect mem write?
	bool FindChainAliasHelper(list<SMPBasicBlock>::iterator CurrBlock, op_t DefOp, ea_t DefAddr);
		// recursive helper for global DefOp with DU-chain that traverses CFG
	bool ConditionalTypePropagation(void); // Apply SCC algorithm to unresolved Phi DEFs
	void MDFindReturnTypes(void); // Fill ReturnRegTypes[]
	void MarkFunctionSafe(void); // Does analysis to see if the function can be marked safe
}; // end class SMPFunction

#endif