#!/usr/bin/python

# This handles the systemtap equivalent of
# $(DTRACE) $(DTRACEFLAGS) -G -s $^ -o $@
# $(DTRACE) $(DTRACEFLAGS) -h -s $^ -o $@
# which is a step that builds DTrace provider and probe definitions

# Copyright (C) 2009 Red Hat Inc.
#
# This file is part of systemtap, and is free software.  You can
# redistribute it and/or modify it under the terms of the GNU General
# Public License (GPL); either version 2, or (at your option) any
# later version.

import os,posix,string,sys
from subprocess import call
from tempfile import mkstemp

class provider:
    semaphores_def = "\n"
    # is the type a basic scalar type?
    def basic_type(self, arg):
        basic_types = [ "int","int*","long","long*","short","short int","unsigned long","char","char*","float","double" ]
        split_arg = arg.rsplit(None,1)
        return (split_arg[0].strip() in basic_types) & (arg.find("[") == -1)
    def typedef_append(self, typedefs,this_probe,arg,c):
        if (add_typedefs):
            split_arg = arg.rsplit(None,1)
            if (self.basic_type(arg)):
                return typedefs
            type_name = " %s_arg%d" % (this_probe.replace("__","_"),c)
            if (len(split_arg) > 1):
                typedefs += ("typedef " + arg.replace(" " + split_arg[1].split("[")[0].lstrip("*"),type_name).strip() + "; ")
                typedefs += (type_name + type_name + "_v;\n")
            else:
                typedefs += ("typedef " + arg.strip() + type_name + "; ")
                typedefs += (type_name + type_name + "_v;\n")
        return typedefs
    def semaphore_def_append(self, this_probe):
        # NB: unsigned short is fixed in ABI
        self.semaphores_def += "__extension__ unsigned short %s_semaphore __attribute__ ((unused)) __attribute__ ((section (\".probes\")));\n" % (this_probe)
    def semaphore_def_write(self, file):
        file.write(self.semaphores_def)
    def generate(self, provider, header, add_typedefs):
        have_provider = False
        self.f = open(provider)
        self.h = open(header,mode='w')
        self.h.write("/* Generated by the Systemtap dtrace wrapper */\n")
        self.h.write("\n#define STAP_HAS_SEMAPHORES 1\n\n")
        self.h.write("\n#include <sys/sdt.h>\n\n")
        in_comment = False
        typedefs = ""
        while (True):
            line = self.f.readline()
            if (line == ""):
                break
            if (line.find("/*") != -1):
                in_comment = True
            if (line.find("*/") != -1):
                in_comment = False
                continue
            if (in_comment):
                continue
            if (line.find("provider") != -1):
                tokens = line.split()
                have_provider = True
                self.provider = tokens[1]
            elif (not have_provider):
                if (add_typedefs):
                    self.h.write (line)
            elif (have_provider and line.find("probe ") != -1):
                while (line.find(")") < 0):
                    line += self.f.readline()
                this_probe = line[line.find("probe ")+5:line.find("(")].strip()
                this_probe_canon = self.provider.upper() + "_" + this_probe.replace("__","_").upper()
                args = (line[line.find("(")+1:line.find(")")])
                args_string = ""
                arg = ""
                i = 0
                c = 0
                self.semaphore_def_append (this_probe)
                while (i < len(args)):
                    if (args[i:i+1] == ","):
                        args_string = ('%s %s,' % (args_string, arg.strip()))
                        c += 1
                        typedefs = self.typedef_append (typedefs,this_probe,arg,c)
                        arg = ""
                    else:
                        arg = arg + args[i]
                    i += 1
                if (i != 0):
                    args_string = ('%s %s' % (args_string, arg.strip()))
                if (len(args_string) == 0):
                    c = 0
                    stap_str = "STAP_PROBE(provider,%s" % (this_probe)
                else:
                    c += 1
                    typedefs = self.typedef_append (typedefs,this_probe,arg,c)
                    stap_str = "STAP_PROBE%d(provider,%s" % (c,this_probe)
                define_str = "#define %s(" % (this_probe_canon)
                i = 1
                while (i <= c):
                    if (i != 1):
                        define_str += ","
                    define_str = define_str + "arg%s" % (i);
                    stap_str = stap_str + ",arg%s" % (i);
                    i += 1
                self.h.write ('/* %s (%s) */\n' % (this_probe_canon,args_string))
                # XXX Enable this when .so semaphores work properly
                self.h.write ('#define %s_ENABLED() %s_semaphore\n' % (this_probe_canon,this_probe))
                # NB: unsigned short is fixed in ABI
                self.h.write ("__extension__ extern unsigned short %s_semaphore __attribute__ ((unused)) __attribute__ ((section (\".probes\")));\n" % (this_probe))
                self.h.write (define_str + ") \\\n")
                self.h.write (stap_str + ")\n\n")
            elif (line.find("}") != -1 and have_provider):
                have_provider = False
        if (add_typedefs):
            self.h.write (typedefs)
        self.h.close()


def usage ():
    print "Usage " + sys.argv[0] + " [--help] [-h | -G] [-C [-I<Path>]] -s File.d [-o <File>]"

def help ():
    usage()
    print "Where -h builds a systemtap header file from the .d file"
    print "      -C when used with -h, also run cpp preprocessor"
    print "      -o specifies an explicit output file name,"
    print "         the default for -G is file.o and -h is file.h"
    print "      -I when running cpp pass through this -I include Path"
    print "      -s specifies the name of the .d input file"
    print "      -G builds a stub file.o from file.d,"
    print "         which is required by some packages that use dtrace."
    sys.exit(1)

def open_file (arg):
    if (len (sys.argv) <= arg):
        return False
    try:
        file = open(sys.argv[arg], 'r')
    except IOError:
        print (sys.argv[arg] + " not found")
        sys.exit(1)
    return file


########################################################################
# main
########################################################################

if (len (sys.argv) < 2):
    usage()
    sys.exit(1)

i = 1
build_header = False
build_source = False
add_typedefs = False
keep_temps = False
use_cpp = False
h_ext = '.h'
filename = ""
s_filename = ""
includes = []
while (i < len (sys.argv)):
    if (sys.argv[i] == "-o"):
        i += 1
        filename = sys.argv[i]
    elif (sys.argv[i] == "-s"):
        i += 1
        s_filename = sys.argv[i]
    elif (sys.argv[i] == "-C"):
        use_cpp = True
    elif (sys.argv[i] == "-h"):
        build_header = True
    elif (sys.argv[i].startswith("-I")):
        includes.append(sys.argv[i])
    elif (sys.argv[i] == "-G"):
        build_source = True
    elif (sys.argv[i] == "-k"):
        keep_temps = True
    elif (sys.argv[i] == "--types"):
        add_typedefs = True
    elif (sys.argv[i] == "--help"):
        help()
    i += 1
if (build_header == False and build_source == False):
    usage()
    sys.exit(1)

if (s_filename != "" and use_cpp):
    (d,fn) = mkstemp(suffix=".d")
    args = ['cpp'] + includes + [s_filename, fn]
    retcode = call(args)
    if (retcode != 0):
        print "\"cpp includes s_filename\" failed"
        usage()
        sys.exit(1)
    s_filename = fn
if (filename == ""):
    if (s_filename != ""):
	(filename,ext) = os.path.splitext(s_filename)
        filename = os.path.basename(filename)
    else:
        usage
        sys.exit(1)
else:
    if (build_header):
        h_ext = ""
    else:
        (filename,ext) = os.path.splitext(filename)
if (build_header):
    providers = provider()
    providers.generate(s_filename, filename + h_ext, add_typedefs)
elif (build_source):
    (basename,ext) = os.path.splitext(s_filename)

    # create for semaphore_def_write
    providers = provider()
    (d,fn) = mkstemp(suffix=".h")
    providers.generate(s_filename, fn, add_typedefs)
    if (not keep_temps):
        os.remove(fn)
    else:
        print "header: " + fn

    (d,fn) = mkstemp(suffix=".c")
    f = open(fn,mode='w')
    f.write("static __dtrace () {}\n")
    providers.semaphore_def_write(f)
    f.close()
    call(["gcc", "-fPIC", "-I.", "-I@prefix@/include", "-g", "-c", fn, "-o", filename + ".o"], shell=False)
    if (not keep_temps):
        os.remove(fn)
    else:
        print "source: " + fn
if (use_cpp):
    if (not keep_temps):
        os.remove(s_filename)
    else:
        print "cpp: " + s_filename