// run.cc (tested on Mac OS X 10.5/10.6)
#include <mach/mach.h>
#include <sys/time.h>
#include <sys/wait.h>
#include <unistd.h>
#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <csignal>

#define SLEEP_NSEC 10000 // 0.01s
// smaller value => fine-grained memory monitoring (overhead increases runtime)

struct mach_send_port_msg {
  mach_msg_header_t          header;
  mach_msg_body_t            body;
  mach_msg_port_descriptor_t task_port;
};

struct mach_recv_port_msg {
  mach_send_port_msg m;
  mach_msg_trailer_t trailer;
};

static int setup_recv_port (mach_port_t *recv_port) {
  mach_port_t   port = MACH_PORT_NULL;
  if (KERN_SUCCESS !=
      mach_port_allocate (mach_task_self (), MACH_PORT_RIGHT_RECEIVE, &port))
    std::fprintf (stderr, "mach_port_allocate failed.\n");
  
  if (KERN_SUCCESS !=
      mach_port_insert_right (mach_task_self (), port, port, MACH_MSG_TYPE_MAKE_SEND))
    std::fprintf (stderr, "mach_port_insert_right failed.\n");
  *recv_port = port;
  return 0;
}

static int send_port (mach_port_t remote_port, mach_port_t port) {
  mach_send_port_msg msg;

  msg.header.msgh_remote_port = remote_port;
  msg.header.msgh_local_port  = MACH_PORT_NULL;
  msg.header.msgh_bits        = MACH_MSGH_BITS (MACH_MSG_TYPE_COPY_SEND, 0) |
                                MACH_MSGH_BITS_COMPLEX;
  msg.header.msgh_size        = sizeof (msg);

  msg.body.msgh_descriptor_count = 1;
  msg.task_port.name             = port;
  msg.task_port.disposition      = MACH_MSG_TYPE_COPY_SEND;
  msg.task_port.type             = MACH_MSG_PORT_DESCRIPTOR;

  if (KERN_SUCCESS != mach_msg_send (&msg.header))
    std::fprintf (stderr, "mach_msg_send failed.\n");

  return 0;
}

static int recv_port (mach_port_t recv_port, mach_port_t *port) {
  mach_recv_port_msg msg;
  
  if (KERN_SUCCESS !=
      mach_msg (&msg.m.header, MACH_RCV_MSG, 0, sizeof (msg), recv_port,
                MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL))
    std::fprintf (stderr, "mach_msg failed.\n");
  
  *port = msg.m.task_port.name;
  return 0;
}

int main (int argc, char **argv) {
  char ** cmd = &argv[1];

  struct timeval start, end;
  task_t         task             = MACH_PORT_NULL;
  mach_port_t    parent_recv_port = MACH_PORT_NULL;
  mach_port_t    child_recv_port  = MACH_PORT_NULL;
  
  // setup port as bootstrap port
  if (setup_recv_port (&parent_recv_port) != 0) return -1;
  if (KERN_SUCCESS !=
      task_set_bootstrap_port (mach_task_self (), parent_recv_port))
    std::fprintf (stderr, "task_set_bootstrap_port failed.\n");
  
  // set timer
  gettimeofday (&start, 0);
  // create new task
  pid_t pid = fork ();
  if (pid < 0)  {
    if (KERN_SUCCESS !=
        mach_port_deallocate (mach_task_self(), parent_recv_port))
      std::fprintf (stderr, "mach_port_deallocate failed.\n");
    std::fprintf (stderr, "cannot fork.\n"); std::exit (1);
  } else if (pid == 0) { // child process
    if (KERN_SUCCESS !=
        task_get_bootstrap_port (mach_task_self (), &parent_recv_port))
      std::fprintf (stderr, "task_get_bootstrap_port failed.\n");
    if (setup_recv_port (&child_recv_port) != 0)              return -1;
    if (send_port (parent_recv_port, mach_task_self ()) != 0) return -1;
    if (send_port (parent_recv_port, child_recv_port)   != 0) return -1;
    if (recv_port (child_recv_port,  &bootstrap_port)   != 0) return -1;
    if (KERN_SUCCESS
        != task_set_bootstrap_port (mach_task_self (), bootstrap_port))
      std::fprintf (stderr, "task_set_bootstrap_port failed.\n");
    execvp (cmd[0], cmd);
    std::fprintf (stderr, "cannot run %s\n", cmd[0]);
    _exit (errno == ENOENT ? 127 : 126);
  }
  // parent process
  std::signal (SIGINT,  SIG_IGN);
  std::signal (SIGQUIT, SIG_IGN);
  if (KERN_SUCCESS !=
      task_set_bootstrap_port (mach_task_self (), bootstrap_port))
    std::fprintf (stderr, "task_set_bootstrap_port failed.\n");
  if (recv_port (parent_recv_port, &task)            != 0) return -1;
  if (recv_port (parent_recv_port, &child_recv_port) != 0) return -1;
  if (send_port (child_recv_port,  bootstrap_port)   != 0) return -1;
  if (KERN_SUCCESS
      != mach_port_deallocate (mach_task_self(), parent_recv_port))
    std::fprintf (stderr, "mach_port_deallocate failed.\n");
  
  // monitor memory
  size_t mem = 0;
  struct task_basic_info t_info;
  mach_msg_type_number_t t_info_count = TASK_BASIC_INFO_COUNT;
  do {
    if (KERN_SUCCESS !=
        task_info (task, TASK_BASIC_INFO, (task_info_t) &t_info, &t_info_count))
      break;
    if (mem < t_info.resident_size) mem = t_info.resident_size;
  } while (! usleep (SLEEP_NSEC));
  
  // stop timer
  int status;
  if (waitpid (pid, &status, 0) == -1) // 
    std::fprintf (stderr, "error waiting for child process.\n");
  if ((status & 0xff) == 0x7f)
    std::fprintf (stderr, "Command stopped by signal %d\n", (status >> 8) & 0xff);
  else if ((status & 0xff) != 0)
    std::fprintf (stderr, "Command terminated by signal %d\n", status & 0xff);
  else if ((status >> 8) & 0xff)
    std::fprintf (stderr, "Command exited with non-zero status %d\n", (status >> 8) & 0xff);
  gettimeofday (&end, 0);

  std::fprintf (stderr, "elapsed (real): %.3fs; RSS=%.2f MiB (%ld bytes)\n",
                end.tv_sec - start.tv_sec + (end.tv_usec - start.tv_usec) * 1e-6,
                mem / 1048576.0, mem);
  std::signal (SIGINT,  SIG_DFL);
  std::signal (SIGQUIT, SIG_DFL);
  return 0;
}