/*
    cpjit.C -- c portable jit in time compiler
    Copyright (C) 2005      Marc Lehmann <cpjit@schmorp.de>
 
    This file is part of libcpjit.

    libcpjit 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 gvpe; if not, write to the Free Software
    Foundation, Inc. 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*/

#include "config.h"

#include "cpjit.h"

#include <cassert>
#include <cstdlib>
#include <cstdarg>
#include <cstdio>
#include <cstring>
#include <cerrno>

#include <iostream>
#include <sstream>
#include <fstream>
#include <stdexcept>
#include <iomanip>

#include <unistd.h>
#include <fcntl.h> // for locking

#include <sys/stat.h>
#include <sys/types.h>
#include <dirent.h>
#include <utime.h>

#include "md5.h"

#include <dlfcn.h>

#define STR_(str) #str
#define STR(str) STR_(str)

#define IDX_VERSION 1

namespace cpjit
{

using namespace std;

const char typestr<U8   >::str[] = "unsigned " STR(TYPE_INT8);
const char typestr<U16  >::str[] = "unsigned " STR(TYPE_INT16);
const char typestr<U32  >::str[] = "unsigned " STR(TYPE_INT32);
const char typestr<U64  >::str[] = "unsigned " STR(TYPE_INT64);
const char typestr<I8   >::str[] = "signed "   STR(TYPE_INT8);
const char typestr<I16  >::str[] =             STR(TYPE_INT16);
const char typestr<I32  >::str[] =             STR(TYPE_INT32);
const char typestr<I64  >::str[] =             STR(TYPE_INT64);
const char typestr<FLT  >::str[] = "float";
const char typestr<DBL  >::str[] = "double";

const char typestr<PTR  >::str[] = "void"                      " *";

const char typestr<U8  *>::str[] = "unsigned " STR(TYPE_INT8)  " *";
const char typestr<U16 *>::str[] = "unsigned " STR(TYPE_INT16) " *";
const char typestr<U32 *>::str[] = "unsigned " STR(TYPE_INT32) " *";
const char typestr<U64 *>::str[] = "unsigned " STR(TYPE_INT64) " *";
const char typestr<I8  *>::str[] = "signed "   STR(TYPE_INT8)  " *";
const char typestr<I16 *>::str[] =             STR(TYPE_INT16) " *";
const char typestr<I32 *>::str[] =             STR(TYPE_INT32) " *";
const char typestr<I64 *>::str[] =             STR(TYPE_INT64) " *";
const char typestr<FLT *>::str[] = "float"                     " *";
const char typestr<DBL *>::str[] = "double"                    " *";

/////////////////////////////////////////////////////////////////////////////

struct digest : md5_ctx
{
  digest ()
  {
    md5_init_ctx (this);
  }

  digest &operator <<(const string &data)
  {
    md5_process_bytes (data.c_str (), data.size (), this);
    return *this;
  }

  operator string();
};

digest::operator string ()
{
  unsigned char data[16];
  std::ostringstream s;

  md5_finish_ctx (this, data);

  for (int i = 0; i < 16; i++)
    s << std::setfill ('0') << std::hex << std::setw (2) << (int)data [i];

  return s.str ();
}

/////////////////////////////////////////////////////////////////////////////

static inline bool
is_null (const string &s)
{
  return s.empty ();
}

/////////////////////////////////////////////////////////////////////////////

// error exception builder
struct error : std::ostringstream
{
  ~error ()
  {
    throw std::runtime_error (str ());
  }
};

#define FATAL(msg) do { error o; o << msg << std::endl; } while (0)
#define ERRNO " (" << strerror (errno) << ")"

/////////////////////////////////////////////////////////////////////////////

static long get_stamp (int fd)
{
  struct stat st;

  if (fstat (fd, &st))
    FATAL ("could not get file stamp");

  return (long)st.st_mtime;
}

static void set_stamp (const string &path, long stamp)
{
  struct utimbuf ut;

  ut.actime  = stamp;
  ut.modtime = stamp;

  if (utime (path.c_str (), &ut))
    FATAL ("could not set file stamp");
}

/////////////////////////////////////////////////////////////////////////////

funsrc::funsrc ()
{
}

void funsrc::update_id (const char *pfx)
{
  digest dgst;

  for (set<string>::const_iterator i = hdrlist.begin ();
       i != hdrlist.end ();
       ++i)
    dgst << *i;

  dgst << retlist << arglist << body;

  id = pfx + (string) dgst;
}

/////////////////////////////////////////////////////////////////////////////

funbuild::funbuild (env &e, const char *rettype, ...)
: e(e)
{
  retlist = rettype;

  std::ostringstream argl;

  va_list ap;
  va_start (ap, rettype);

  int args = 0;
  while ((rettype = va_arg (ap, const char *)))
    {
      if (args++)
        argl << ", ";

      argl << rettype << " a" << args;
    }

  va_end (ap);

  arglist = argl.str ();
}

void funbuild::header (const std::string &hdr)
{
  hdrlist.insert (hdr);
}

funbuild &funbuild::operator <<(const char *src)
{
  bodystream << src;
}

funbuild &funbuild::operator <<(const string &src)
{
  bodystream << src;
}

fun *funbuild::get ()
{
  body = bodystream.str ();
  update_id ("f_");
  e.register_fragment (*this);

  return new fun (e, id);
}

/////////////////////////////////////////////////////////////////////////////

fun::fun (env &e, const funid &id)
: e(e), id(id), funptr(0)
{
}

void *fun::ptr ()
{
  if (!funptr)
    funptr = e.sym (id, id.c_str ());

  return funptr;
}

/////////////////////////////////////////////////////////////////////////////

env::env (const string &path, bool temporary)
: path(path), temporary(temporary)
{
  bool create;
  mode_t oflags = O_RDWR | O_EXCL;

  if (is_null (path))
    {
      temporary = true;
      this->path = tmpnam (0);
    }

  create = mkdir (this->path.c_str (), 0777) == 0;

  if (!create && temporary)
    FATAL ("error while creating libcpjit temporary directory" << ERRNO);
  
  idxpath = this->path + "/libcpjit.idx";
  idxfd = open (idxpath.c_str (), O_RDWR | (create ? O_CREAT | O_EXCL : 0), 0666);
  if (idxfd < 0)
    FATAL (path << ": unable to open index file " << idxpath);

  lock ();

  if (create)
    {
      idxstamp = get_stamp (idxfd) - 1;
      nextid = 0;
      save_index ();
    }

  load_index ();
  unlock ();
  compact ();
}

env::~env ()
{
  if (temporary)
    clear ();
}

void env::clear ()
{
  DIR *dir = opendir (path.c_str ());

  if (dir)
    {
      struct dirent *de;
      while ((de = readdir (dir)))
        {
          if (de->d_name [0] == '.')
            continue;

          string f = path + "/" + de->d_name;
          unlink (f.c_str ());
        }

      closedir (dir);
    }

  if (rmdir (path.c_str ()))
    FATAL (path << ": unable to remove libcpjit directory" << ERRNO);
}

void env::dolock (bool lock)
{
  struct flock fl;

  fl.l_type   = lock ? F_WRLCK : F_UNLCK;
  fl.l_whence = SEEK_SET;
  fl.l_start  = 0;
  fl.l_len    = 0;

  while (fcntl (idxfd, F_SETLKW, &fl) < 0)
    if (errno != EINTR)
      FATAL (path << ": unable to lock index file" << ERRNO);
}

void env::lock ()
{
  dolock (true);
}

void env::unlock ()
{
  dolock (false);
}

bool
env::load_index ()
{
  // caller must lock()

  long stamp = get_stamp (idxfd);

  if (stamp == idxstamp)
    return false;

  idxstamp = stamp;

  fstream f (idxpath.c_str (), ios::in);

  if (!f)
    FATAL (idxpath << ": unable to read index" << ERRNO);

  f >> nextid;

  if (nextid != IDX_VERSION)
    FATAL (idxpath << ": index file corrupted");

  f >> nextid;

  id2file.clear ();

  while (f)
    {
      string id; f >> id;

      if (id == "-")
        break;

      string fileid; f >> fileid;

      id2file [id] = fileid;

    }

  string end; f >> end;

  if (end != "__end__" || !f)
    FATAL (idxpath << ": index file corrupted");

  f.close ();

  return true;
}

void
env::save_index ()
{
  // caller must lock()+load_index()

  std::fstream f (idxpath.c_str ());

  if (!f)
    FATAL (idxpath << ": unable to write index" << ERRNO);

  f << IDX_VERSION << "\n"
    << nextid << "\n";

  for (map<string, string>::const_iterator i = id2file.begin ();
       i != id2file.end ();
       ++i)
    f << (*i).first << " " << (*i).second << "\n";

  f << "-\n";
  f << "__end__\n";

  f.close ();

  set_stamp (idxpath, ++idxstamp);
}

void
env::compact ()
{
  lock ();
  save_index ();
  unlock ();
}

void env::register_fragment (const funsrc &fragment)
{
  fragments.push_back (fragment);
}

void *env::dlsym (const string &id, const char *symbol)
{
  assert (id2file.find (id) != id2file.end ());

  const string &fileid = id2file [id];

  void *so;

  // load so
  if (file2so.find (fileid) == file2so.end ())
    {
      string sopath = path + "/" + fileid + ".so";

      so = dlopen (sopath.c_str (), RTLD_LAZY | RTLD_GLOBAL);
      if (!so)
        FATAL (sopath << ": so file should be there but isn't" << ERRNO);
    }
  else
    so = file2so [fileid];

  return ::dlsym (so, symbol);
}

void *env::sym (const string &id, const char *symbol)
{
  // not yet compiled?
  if (id2file.find (id) == id2file.end ())
    {
      lock ();

      // nobody else compiled it?
      if (!load_index () || id2file.find (id) == id2file.end ())
        {
          ostringstream sfile;
          sfile << nextid++;
          string fileid = sfile.str ();
          string base = path + "/" + fileid;

          string fc  = base + ".c";
          string fo  = base + ".o";
          string fso = base + ".so";

          std::ofstream f (fc.c_str ());
          set<string> hdrs;

          for (vector<funsrc>::const_iterator i = fragments.begin ();
               i != fragments.end ();
               ++i)
            if (id2file.find (i->id) == id2file.end ())
              {
                f << "#line 1 \"" << i->id << "-hdrlist\"\n";

                for (set<string>::const_iterator j = i->hdrlist.begin ();
                     j != i->hdrlist.end ();
                     ++j)
                  f << *j << "\n";

                f << "#line 1 \"" << i->id << "-body\"\n"
                  << i->retlist << " " << i->id << " (" << i->arglist << ")\n"
                  << "{\n"
                  << i->body
                  << "\n}\n";

                id2file [i->id] = fileid;
              }
          
          f.close ();
          fragments.clear ();

          system ((CC " " CFLAGS " -o " + fo + " " + fc).c_str ());
          unlink (fc.c_str ());
          system ((LD " " LFLAGS " -o " + fso + " " + fo).c_str ());

          save_index ();
        }

      unlock ();
    }

  void *p = dlsym (id, symbol);

  if (!p)
    {
      lock ();
      load_index ();
      unlock ();
      p = dlsym (id, symbol);
    }

  if (!p)
    FATAL ("symbol " << id << ":" << symbol << " not where it should be");

  return p;
}

}