#!/bin/bash 

LANG="C"
export LANG

eval `tau-config`

if [ "x$MAQAO_BINARY" = "x" ]; then
  maqao_bin=$BASEDIR/bin/tau_maqao
else
  maqao_bin=$MAQAO_BINARY
  echo "tau_rewrite: Using maqao binary from MAQAO_BINARY environment variable: $maqao_bin"
fi

if [ ! -x $maqao_bin ] 
then
  echo "$0: Couldn't locate smaqao in PDT"
  if [ -x $BASEDIR/bin/tau_run ]
  then 
    echo "$0: Using DyinstAPI for binary rewriting" 
    arglist=""
    lastarg=""
    output_file_specified=false
    for arg in "$@"
    do 
      arglist="$arglist $lastarg"
      lastarg=$arg
      if [ $arg = "-o" ]
      then
        output_file_specified=true
      fi
    done
    if [ $output_file_specified = false ] 
    then
      arglist="$arglist -o"
    fi
    arglist="$arglist $lastarg" 
    $BASEDIR/bin/tau_run $arglist
    exit 1;
  else
    echo "$0: To use TAU's binary rewriting capabilities, please install MAQAO from PDT v3.17+ (x86_64 only at present) or DyninstAPI"
    exit 1; 
  fi
fi


echoIfVerbose () {
    if [ $verbose = "true" ] ; then
	echo -e "$1"
    fi
}

echoIfShortVerbose () {
    if [ $shortverbose = "true" ] ; then
	echo -e "$1"
    fi
}

usage()
{
    echo ""
    echo "Usage: tau_rewrite [options] [--] <exe> <exe options>"
    echo ""
# Common options first
    echo "Options:"

    echo "        -o <outfile>       : specify instrumented output file"
    options=`tau-config --list-options`
    echo "        -T <$options> : specify TAU option"
    echo "        -loadlib=<file.so>   : specify additional load library"
    echo "        -s : Dryrun without executing"
    echo "        -v       : long verbose mode"
    echo "        -v1       : short verbose mode"
    echo "        -XrunTAUsh-<options> : specify TAU library directly"
    echo ""
    echo "Notes:"
    echo "	Defaults if unspecified: -T MPI"
    echo "	MPI is assumed unless SERIAL is specified"
    echo ""
    echo "Example:"
    echo "    tau_rewrite -T papi,pdt a.out -o a.inst"
    echo "    mpirun -np 4 ./a.inst"
    echo ""
    exit
}

if [ $# = 0 ] ; then
    usage
fi


dryrun=false
processT=false
TauOptions=""
TauOptionsExclude=""
verbose=false
shortverbose=false
binding_specified=""
binding_options=""
extraloadlibs=""
scorep=false
tauluafile="tau_rewrite.lua"
processf="false"
includelist="false"
looplist="false"
excludelist="false"

for arg in "$@" ; do
  # Thanks to Bernd Mohr for the following that handles quotes and spaces (see configure for explanation)
  modarg=`echo "x$arg" | sed -e 's/^x//' -e 's/"/\\\"/g' -e s,\',%@%\',g -e 's/%@%/\\\/g' -e 's/ /\\\ /g'`

  if [ "$processT" = true ] ; then
      binding_options=`echo $binding_options $arg | sed -e 's/,/ /g' | tr '[A-Z]' '[a-z]'`
      processT="false"
      test_arg=`echo $arg | sed -e 's@scorep@@g'`
      if [ "x$test_arg" != "x$arg" ]; then
        scorep=true
      fi
      shift
  else
    if [ "$processf" = true ] 
    then
      selectfile=$arg
      processf="false"
      shift
    else
      case $arg in 
	  -v|-d|-verbose|--verbose)
	      verbose=true
	      shift
	      ;;
	  -v1)
	      shortverbose=true
	      shift
	      ;;
	  -h|-help|--help)
	      usage
	      ;;
	  -s)
	      dryrun=true
	      shift
	      ;;
	  -V)
	      echo '$Id: tau_exec,v 1.19 2010/06/09 18:11:25 amorris Exp $';
	      exit 0;
	      ;;
	  -T)
	      processT=true
	      shift
	      ;;
	  -tau:*)
              binding_options="$binding_options `echo $arg | sed -e 's/-tau://' -e 's/,/ /g'`"
	      shift
              ;;
	  -f)
          processf=true
          echoIfShortVerbose "Process F is true!"
          shift;;
	  -loadlib=*)
	      myarg=`echo $arg | sed 's/-loadlib=//'`
              if [ "x$extraloadlibs" = "x" ]
              then
                extraloadlibs="$myarg"
              else
                extraloadlibs="$myarg $extraloadlibs"
              fi
	      shift
              ;;
	  -XrunTAU-*)
	      myarg=`echo $arg | sed 's/-XrunTAU-//'`
	      binding_specified="shared-$myarg"
	      shift
	      ;;
	  -XrunTAUsh-*)
	      myarg=`echo $arg | sed 's/-XrunTAUsh-//'`
	      binding_specified="shared-$myarg"
	      shift
	      ;;
	  --)
	      shift
	      break
	      ;;
	  -*)
	      echo "Unknown option: $arg" >&2
	      exit 1
# First non-option signifies end of options. This would be much easier with getopt()
	      ;;
	  *)
	      break
	      ;;
      esac
    fi
  fi
done


if [ "x$selectfile" != "x" ] 
then
  echoIfShortVerbose "selectfile=$selectfile"
  if [ `grep "^loops" $selectfile  | wc -l ` != 0 ] 
  then
    looplist="true"
    echoIfShortVerbose "looplist=$looplist"
  fi
  if [ `grep "^BEGIN_EXCLUDE_LIST" $selectfile  | wc -l ` != 0 ] 
  then
    excludelist="true"
    echoIfShortVerbose "excludelist=$excludelist"
  fi
  if [ `grep "^BEGIN_INCLUDE_LIST" $selectfile  | wc -l ` != 0 ] 
  then
    includelist="true"
    echoIfShortVerbose "includelist=$includelist"
  fi
fi

# choose TAU library
new_binding_options=""
if [ "x$binding_options" != "x" ]; then
    for i in $binding_options ; do
      case $i in 
	  *)
	      new_binding_options="$new_binding_options $i"
	      ;;
      esac
    done
fi
binding_options="$new_binding_options"


if [ "x$binding_specified" = "x" ] ; then
    if [ "x$binding_options" = "x" ]; then
	binding_options=$DEFAULT_BINDING
    else
        # Add MPI by default
	serial=`echo $binding_options | grep serial`
	if [ $? != 0 ] ; then
           # add mpi if shmem is not specified
	    shmem=`echo $binding_options | grep shmem`
	    if [ $? != 0 ] ; then
	        binding_options="$binding_options mpi"
            fi
	fi
    fi
    theBinding=`tau-config --binding $binding_options`
    if [ $? != 0 ] ; then
	exit 1
    fi
else
    theBinding=$binding_specified
fi



    echoIfVerbose ""
    echoIfVerbose "Program to run : $@"
    echoIfVerbose ""

if [ `uname -s ` = Darwin ]; then
  apple=1
  TAU_SHLIBX=.dylib
else
  apple=0
  TAU_SHLIBX=.so
fi

TAUEX_LD_LIBRARY_PATH=$BASEDIR/lib/$theBinding
TAU_INST_LIBRARY=$BASEDIR/lib/$theBinding/libTAU$TAU_SHLIBX
TAU_MOD_INST_LIBRARY=$TAU_INST_LIBRARY

if [ $scorep = true ]; then
  scoreplibrary=`ldd $BASEDIR/lib/$theBinding/libTAU.so | grep scorep_adapter_mpi_event.so | awk ' { print $3;}'`
  scorepdir=`echo $scoreplibrary | sed -e 's@libscorep_adapter_mpi_event.so.0@@g'`
  echoIfVerbose "scorepdir= $scorepdir"
  if [ -f $scoreplibrary -a -d $scorepdir ]; then 
    echoIfVerbose "Preloading: $scoreplibrary from $scorepdir"
    TAU_MOD_INST_LIBRARY="$BASEDIR/lib/$theBinding/libTAU$TAU_SHLIBX:$scoreplibrary:$BASEDIR/lib/$theBinding/libTAU-preload.so"
# NOTE: We should put this list of libraries in the LD_PRELOAD in $outfile but not inject it in the binary. 
    TAU_MOD_LD_LIBRARY_PATH="$TAUEX_LD_LIBRARY_PATH:$scorepdir"
    echoIfVerbose "TAU_MOD_LD_LIBRARY_PATH: $TAU_MOD_LD_LIBRARY_PATH"
  fi
fi

# add libraries specified by -loadlib=<foo.so>
#TAU_INST_LIBRARY=${TAU_INST_LIBRARY}${extraloadlibs}

# remove double colons
#TAU_INST_LIBRARY=`echo $TAU_INST_LIBRARY | sed -e "s/::/:/g" -e "s/:$//"`

if [ $verbose = "true" ] ; then
    echo "Matching bindings:"
    tau-config --list-matching $binding_options
    echo ""
    echo "Using:"
    echo "$theBinding"
    echo ""
    echo "Configuration:"
    echo ""
    echo "Setting LD_LIBRARY_PATH to $TAUEX_LD_LIBRARY_PATH"
    echo "Using instrumentation library $TAU_INST_LIBRARY"
    echo ""
fi

echoIfShortVerbose "Using: $TAU_INST_LIBRARY"
echoIfShortVerbose "Additional libs = ${extraloadlibs}"

processOutputFile=false
# First argument is the executable
infile="$arg"
for arg in "$@" ; do
  if [ "$processOutputFile" = true ]
  then 
    outfile=$arg
  else  
  echoIfShortVerbose "arg  = $arg"
  case $arg in 
    -o) processOutputFile="true"
	shift;;
    *) 
	shift;;
  esac;
  fi
done

echoIfShortVerbose "infile=$infile"
if [ "x$outfile" = "x" ] 
then
  echoIfShortVerbose "infile=$infile, arg = $arg"
  if [ "x$infile" = "x$arg" ] 
  then 
    outfile="$infile.tau_inst"
    echo "$0: Output file written to $outfile"
  else
    outfile=$arg
  fi
fi

echoIfShortVerbose "outfile=$outfile"

if [ "$dryrun" = true ]
then
  echoIfShortVerbose "Dryrun"
  exit 1;
fi

outputdir=`dirname ${infile}`
echoIfShortVerbose "dirname = $outputdir"
baseinfile=`basename ${infile}`
if [ "$outputdir" = "." ]
then
  outputdir=`pwd` 
fi

# define functions in the lua file
defs="function isoutermost(loop)
  if (loop:is_outermost() == true) then
    return true;
  end 
end 
          
function getloopstring(loop) 
  local filename = loop:get_function():get_src_file_name();
  local functionname;
  local startloop, stoploop = loop:get_src_line();
  local returnstring;
            
  if(loop:get_function():get_demname()~= nil)
  then
    functionname = loop:get_function():get_demname();
  else  
    functionname = loop:get_function():get_name();
  end     
  returnstring=\"Loop: \"..functionname..\" [{\"..filename..\"} {\"..startloop..\",0}-{\"..stoploop..\"}]\";
  print (returnstring); 
  return returnstring;
end 
" 

# Get the list of external libs
prologue=""
makefile_binding=`echo $theBinding | sed -e 's@shared@@g'`


if [ $apple = 1 ]; then
  TAU_LDD=`otool -L $infile`
  isStatic=$?
else
  if [ "x$TAUARCH" = "xmic_linux" ]; then
    # Does the maqao binary support --check-dynamic flag?
    if [ `$maqao_bin madras --check-dynamic ./foo 2>&1  | grep unrecognized | wc -l ` = 0 ]; then
      echoIfShortVerbose "$maqao_bin supports --check-dynamic"
      TAU_LDD=`$maqao_bin madras  --check-dynamic $infile 2>&1 | grep "is not a dynamic executable" | wc -l `
      if [ $? = 0 ]; then
        isStatic=0
      fi
     else
      TAU_LDD=`$maqao_bin madras  --get-dynamic-lib $infile | grep 'is not a dynamic executable' | wc -l `
      if [ $? = 0 ]; then
        isStatic=1
      fi
    fi
    echo "Using ldd=$TAU_LDD"
  else
    TAU_LDD=`ldd $infile`
    isStatic=$?
  fi
fi



if [ $isStatic = 1 ]; then
  echoIfShortVerbose "STATIC EXECUTABLE $infile"
  isStatic=true
  TAU_INST_LIBRARY=$BASEDIR/lib/libtau$makefile_binding.a
  echoIfShortVerbose "Using: $TAU_INST_LIBRARY"
else
  echoIfShortVerbose "Dynamic executable $infile"
fi

if [ $isStatic = 0 ]; then
  prologue=""
else
  prologue="extlibs = { 
`eval $BASEDIR/bin/tau_show_libs $BASEDIR/lib/Makefile.tau$makefile_binding`  
};"
fi

# Process loop text if it is relevant
if [ $looplist = true ]
then
  looproutines=`sed '/BEGIN_INSTRUMENT_SECTION/,/END_INSTRUMENT_SECTION/{/BEGIN_INSTRUMENT_SECTION/{h;d};H;/END_INSTRUMENT_SECTION/{x;/BEGIN_INSTRUMENT_SECTION/,/END_INSTRUMENT_SECTION/p}};d' $selectfile | grep loops | grep routine | sed -e 's/loops\( *\)//' -e 's/routine=//' -e 's/#/\.\*/g'  | sed -n '1h;2,$H;${g;s/\n/,/g;p}' `
  echoIfShortVerbose "Loop routines=$looproutines"
  looptext=", {
          filters = {{  
            type = \"whitelist\", 
            filter = { {subtype = \"regexplist\",value = {$looproutines}} }
          }}, 
          loops = {{
            filters={{
                type = \"user\",
                filter = isoutermost;
            }},
            entries = {{
                 at_program_entry = { 
                   {   
                      name = \"tau_trace_register_loop\",
                      lib = \"$TAU_INST_LIBRARY\", 
                      params = { 
                        {type = \"macro\",value = \"profiler_id\"}, 
                        {type = \"function\",value = getloopstring},
                      }   
                   }   
                 },  
                 name = \"tau_loop_trace_entry\",
                 lib = \"$TAU_INST_LIBRARY\", 
                 params = { 
                     {type = \"macro\",value = \"profiler_id\"}
                   }   
                 }}, 
            exits = {{
               name = \"tau_loop_trace_exit\",
               lib = \"$TAU_INST_LIBRARY\", 
               params = { 
                 {type = \"macro\",value = \"profiler_id\"}
               }   
            }}  
          }}, 
         }, 
  " 
fi

if [ $includelist = false ] 
then 
  # Process exclude text if it is relevant
  if [ $excludelist = true ]
  then
    excluderoutines=` sed '/BEGIN_EXCLUDE_LIST/,/END_EXCLUDE_LIST/{/BEGIN_EXCLUDE_LIST/{h;d};H;/END_EXCLUDE_LIST/{x;/BEGIN_EXCLUDE_LIST/,/END_EXCLUDE_LIST/p}};d' $selectfile |  sed -e 's/BEGIN_EXCLUDE_LIST//' -e 's/END_EXCLUDE_LIST//' -e 's/#/\.\*/g'  -e 's/"//g' -e 's/^/"/' -e 's/$/"/' | sed -n '1h;2,$H;${g;s/\n/,/g;p}' | sed -e 's/"",//g' -e 's/,""//g' `
  
    echoIfShortVerbose "Include Routines are $excluderoutines"
    excludetext=" 
          filters = {{  
            type = \"blacklist\", 
            filter = { {subtype = \"regexplist\",value = {\"_init\", $excluderoutines}} }
            }},
    "
  else
    excludetext=" 
          filters = {{  
            type = \"blacklist\", 
            filter = { {subtype = \"regexplist\",value = {\"_init\"}} },
          {subtype = \"stringlist\",value = {\"call_gmon_start\",\"_init\",\"_fini\",\"_start\",\"frame_dummy\"}},
          {subtype = \"regexplist\",value = {\"^irc__.*\",\"^__do_global_.*\",\"^__libc_csu_.*\",\"^__intel_.*\",\"^__pthread_.*\",\"pgf90_.*\", \"pghpf_\", \"__hpf_\"}}
            }},
    "
  fi
fi

# Process include text if it is relevant
if [ $includelist = true ]
then
  includeroutines=`sed '/BEGIN_INCLUDE_LIST/,/END_INCLUDE_LIST/{/BEGIN_INCLUDE_LIST/{h;d};H;/END_INCLUDE_LIST/{x;/BEGIN_INCLUDE_LIST/,/END_INCLUDE_LIST/p}};d' $selectfile |  sed -e 's/BEGIN_INCLUDE_LIST//' -e 's/END_INCLUDE_LIST//' -e 's/#/\.\*/g'  -e 's/"//g' -e 's/^/"/' -e 's/$/"/' | sed -n '1h;2,$H;${g;s/\n/,/g;p}' | sed -e 's/"",//g' -e 's/,""//g'`
  echoIfShortVerbose "Include Routines are $includeroutines"
  includetext=" 
          filters = {{  
            type = \"whitelist\", 
            filter = { {subtype = \"regexplist\",value = {$includeroutines}} }
            }},
  "
fi
# Now we write to the $outfile
luafile="$defs $prologue
events = {
  run_dir = \"$outputdir/\",
  functions_global_blacklist = {
          {subtype = \"stringlist\",value = {\"call_gmon_start\",\"_init\",\"_fini\",\"_start\",\"frame_dummy\"}},
          {subtype = \"regexplist\",value = {\"^irc__.*\",\"^__do_global_.*\",\"^__libc_csu_.*\",\"^__intel_.*\",\"^__pthread_.*\",\"pgf90_.*\", \"pghpf_\", \"__hpf_\"}}
  },
  at_exit={
          {
            name = \"tau_dyninst_cleanup\",
            lib = \"$TAU_INST_LIBRARY\"
          }
  },
  main_bin = {
        properties={
          enable_function_instrumentation = true, 
          distinguish_inlined_functions = true,
          distinguish_suffix = \"_omp\", 
          generate_runfile   = true
        },
        path= \"$outputdir/$baseinfile\",
        output_suffix = \"_inst\",
        output_name = \"$baseinfile.tau_rewritten\",
        functions={
          { $includetext $excludetext
            entries = {
              {
                 at_program_entry = {
                   {
                      name = \"trace_register_func\",
                      lib = \"$TAU_INST_LIBRARY\",
                      params = {
                        {type = \"macro\",value = \"fct_info_summary\"},
                        {type = \"macro\",value = \"profiler_id\"},
                      }
                   }
                 },
                 name = \"traceEntry\",
                 lib = \"$TAU_INST_LIBRARY\",   
                 params = {
                            {type = \"macro\",value = \"profiler_id\"}
                 }
              }
            },
            exits = {
              {
                 name = \"traceExit\",
                 lib = \"$TAU_INST_LIBRARY\",   
                 params = {
                   {type = \"macro\",value = \"profiler_id\"}
                 }
              }
            }
          } $looptext
        }
      }"
echoIfVerbose "$luafile"
echo "$luafile"> $tauluafile

for extralibs in ${extraloadlibs}
do
  luafile=",  libs = {
        {
                properties={
                  enable_function_instrumentation = true
                  distinguish_inlined_functions = true,
                  distinguish_suffix = \"_omp\", 
                },
                path = \"$extralibs\",
                output_suffix = \"_inst\",
                output_name = \"lib_inst.so\",
                functions={
                  {
                    entries = {
                      {
                         at_program_entry = {
                           {
                              name = \"trace_register_func\",
                              lib = \"$TAU_INST_LIBRARY\",
                              params = {
                                {type = \"macro\",value = \"fct_info_summary\"},
                                {type = \"macro\",value = \"profiler_id\"},
                              }
                           }
                         },
                         name = \"traceEntry\",
                         lib = \"$TAU_INST_LIBRARY\",   
                         params = {
                            {type = \"macro\",value = \"profiler_id\"}
                         }
                      }
                    },
                    exits = {
                      {
                         name = \"traceExit\",
                         lib = \"$TAU_INST_LIBRARY\",   
                         params = {
                           {type = \"macro\",value = \"profiler_id\"}
                         }
                      }
                    }
                  }
                }

        }
  }"
  echoIfVerbose "$luafile"
  echo "$luafile" >> $tauluafile
done
echoIfVerbose "};"
echo "};" >> $tauluafile

echoIfShortVerbose " Output written!"

cat << EOF > tau_empty.in


EOF


if [ "$dryrun" = true ]
then
  echo "smaqao module=mil input=tau_rewrite.lua"
else
  eval `$maqao_bin module=mil input=tau_rewrite.lua <tau_empty.in &> tau_instr.log `
  echo "tau_rewrite: Binary instrumentation done through MAQAO Multi-Architecture Disassembler, Rewriter and ASsembler technology"
  if [ $verbose = true ] 
  then
    cat tau_instr.log
  fi
# Move the output to the outputfile. 
  ls -t *_0.run &> /dev/null
  if [ $? = 1 ]
  then
    echo "TAU: There was a problem instrumenting the executable using MAQAO. Please rerun the command with the -v option for further details"
  else
    mv  `ls -t *_0.run | head -1 ` $outfile~
    echo "#!/bin/sh
          if test \"x\$LD_LIBRARY_PATH\" = \"x\" ; then
            export LD_LIBRARY_PATH=$TAU_MOD_LD_LIBRARY_PATH
          else 
            export LD_LIBRARY_PATH=$TAU_MOD_LD_LIBRARY_PATH:\$LD_LIBRARY_PATH
          fi

          if test \"x\$LD_PRELOAD\" = \"x\" ; then 
             export LD_PRELOAD=$TAU_MOD_INST_LIBRARY
          else
             export LD_PRELOAD=$TAU_MOD_INST_LIBRARY:\$LD_PRELOAD
          fi" > $outfile_1~
    cat $outfile_1~ $outfile~ > $outfile
    chmod +x $outfile
    /bin/rm -f $outfile_1~ $outfile~
  fi
  if [ $shortverbose = false -a $verbose = false ]
  then
    echo ""
# Not executing /bin/rm -f instru_sessions.lua madras_trace.log tau_instr.log tau_rewrite.lua
  else 
      /bin/rm -f tau_instr.log tau_empty.in tau_rewrite.lua
  fi
fi
