#!/usr/bin/python
# vim: et sta sts=4 sw=4 ts=8

# 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, 2010, 2011 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
import posix
import sys
import shlex
from subprocess import call
from tempfile import mkstemp

class _provider(object):
    def __init__(self):
        self.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, add_typedefs):
        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 += '#if defined STAP_SDT_V1\n'
        self.semaphores_def += '#define %s_%s_semaphore %s_semaphore\n' % \
            (self.provider,this_probe,this_probe)
        self.semaphores_def += '#endif\n'
        self.semaphores_def += '#if defined STAP_SDT_V1 || defined STAP_SDT_V2 \n'
        self.semaphores_def += "__extension__ unsigned short %s_%s_semaphore __attribute__ ((unused)) __attribute__ ((section (\".probes\")));\n" % \
            (self.provider,this_probe)
        self.semaphores_def += '#else\n'
        self.semaphores_def += "__extension__ unsigned short %s_%s_semaphore __attribute__ ((unused)) __attribute__ ((section (\".probes\"))) __attribute__ ((visibility (\"hidden\")));\n" % \
            (self.provider,this_probe)
        self.semaphores_def += '#endif\n'

    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 _SDT_HAS_SEMAPHORES 1\n\n")
        self.h.write("\n#define STAP_HAS_SEMAPHORES 1 /* deprecated */\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, add_typedefs)
                        arg = ""
                    else:
                        arg = arg + args[i]
                    i += 1
                if (i != 0):
                    args_string = ('%s %s' % (args_string, arg.strip()))
                if (not args_string):
                    c = 0
                    stap_str = "DTRACE_PROBE(%s,%s" % \
                    (self.provider,this_probe)
                else:
                    c += 1
                    typedefs = self.__typedef_append (typedefs, this_probe,
                                                      arg, c, add_typedefs)
                    stap_str = "DTRACE_PROBE%d(%s,%s" % \
                        (c,self.provider,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))
                self.h.write('#if defined STAP_SDT_V1\n')
                self.h.write('#define %s_ENABLED() __builtin_expect (%s_semaphore, 0)\n' % \
                                  (this_probe_canon,this_probe))
                self.h.write('#define %s_%s_semaphore %s_semaphore\n' % \
                                  (self.provider,this_probe,this_probe))
                self.h.write('#else\n')
                self.h.write('#define %s_ENABLED() __builtin_expect (%s_%s_semaphore, 0)\n' % \
                                  (this_probe_canon,self.provider,this_probe))
                self.h.write('#endif\n')
                # NB: unsigned short is fixed in ABI
                self.h.write("__extension__ extern unsigned short %s_%s_semaphore __attribute__ ((unused)) __attribute__ ((section (\".probes\")));\n" % \
                                  (self.provider,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)


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

def main():
    if (len(sys.argv) < 2):
        usage()
        return 1

    i = 1
    build_header = False
    build_source = False
    add_typedefs = False
    keep_temps = False
    use_cpp = False
    suffix = ""
    filename = ""
    s_filename = ""
    includes = []
    defines = []
    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].startswith("-D")):
            defines.append(sys.argv[i])
        elif (sys.argv[i] == "-h"):
            build_header = True
            suffix = ".h"
        elif (sys.argv[i].startswith("-I")):
            includes.append(sys.argv[i])
        elif (sys.argv[i] == "-G"):
            build_source = True
            suffix = ".o"
        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()
        return 1

    if (s_filename != "" and use_cpp):
        (d,fn) = mkstemp(suffix=".d")
        CPP = os.environ.get("CPP", "cpp")
        retcode = call(shlex.split(CPP) + includes + defines + [s_filename, fn])
        if (retcode != 0):
            print "\"cpp includes s_filename\" failed"
            usage()
            return 1
        s_filename = fn
    if (filename == ""):
        if (s_filename != ""):
            (filename,ext) = os.path.splitext(s_filename)
            filename = os.path.basename(filename)
        else:
            usage()
            return 1
    else:
        suffix = ""
    if (build_header):
        providers = _provider()
        providers.generate(s_filename, filename + suffix, 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')
        # dummy declaration just to make the object file non-empty
        f.write("static void __dtrace (void) __attribute__((unused));\n")
        f.write("static void __dtrace (void) {}\n")
        f.write("\n#include <sys/sdt.h>\n\n")
        providers.semaphore_def_write(f)
        f.close()
        CC = os.environ.get("CC", "gcc")
        CFLAGS = "-g " + os.environ.get("CFLAGS", "")
        retcode = call(shlex.split(CC) + defines + includes + shlex.split(CFLAGS) +
             ["-fPIC", "-I.", "-I/usr/include", "-c", fn, "-o",
              filename + suffix], shell=False)
        if (retcode != 0):
            print "\"gcc " + fn + "\" failed"
            usage()
            return 1
        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
    return 0

if __name__ == "__main__":
    sys.exit(main())