mtsread/mtsread.c

/*  mtsread - Minetest schematic reader.
    Copyright (C) 2022, 2023 Noisytoot

    This program 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 3 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 this program.  If not, see <https://www.gnu.org/licenses/>.  */

#include <stdbool.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <inttypes.h>
#include <fcntl.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <string.h>
#include <zlib.h>

/* Constants */
#define MTSCHEM_FILE_SIGNATURE  0x4d54534d /* 'MTSM' */
#define MTSCHEM_FILE_VERSION    4

#define BUFSIZE 512

enum image_format {
  IMGFMT_PPM_TEXT,
  IMGFMT_PPM_BINARY,
  IMGFMT_PAM,
};

#define is_newline(c) ((c) == '\n' || (c) == '\r')
#define is_whitespace(c) ((c) == '\t' || (c) == ' ')

#define lengthof(array) (sizeof(array)/sizeof(array[0]))

typedef struct mts_header {
  uint32_t signature;
  uint16_t version;
  uint16_t size_x;
  uint16_t size_y;
  uint16_t size_z;
  uint8_t *slice_probs;
  uint16_t name_id_count;
  uint8_t *name_id_map;
} mts_header;

typedef struct mts {
  mts_header header;
  uint16_t  *param0; /* Node IDs (length 2*X*Y*Z, offset 0) */
  uint8_t   *param1; /* Probabilities (length X*Y*Z, offset 2*X*Y*Z) */
  uint8_t   *param2; /* param2 (length X*Y*Z, offset 3*X*Y*Z) */
} mts;

typedef struct colour {
  uint8_t red;
  uint8_t green;
  uint8_t blue;
  uint8_t alpha;
} colour;

typedef struct image {
  uint16_t width;
  uint16_t height;
  uint16_t colours;
  /* Indexed by colourmap */
  uint16_t *pixmap;
  colour *colourmap;
} image;

#define foreach_node(schematic, x, y, z)                         \
  for (uint16_t x = 0; x < schematic->header.size_x; x++)        \
    for (uint16_t y = 0; y < schematic->header.size_y; y++)      \
      for (uint16_t z = 0; z < schematic->header.size_z; z++)

ssize_t read_or_fail(int fd, void *buf, size_t count) {
  ssize_t result = read(fd, buf, count);
  if (result < (ssize_t)count) {
    fprintf(stderr, "Unexpected EOF\n");
    exit(1);
  }
  return result;
}

void *calloc_or_fail(size_t nmemb, size_t size) {
  void *buf = calloc(nmemb, size);
  if (!buf) {
    fprintf(stderr, "Failed to allocate memory\n");
    exit(1);
  }
  return buf;
}

void read_header(int fd, mts_header *header) {
  char buf[UINT16_MAX];
  off_t offset;

  /* Read first part of header */
  read_or_fail(fd, &buf, 12);
  header->signature = ntohl(*(uint32_t *)buf);
  header->version = ntohs(*(uint16_t *)(buf + 4));
  header->size_x = ntohs(*(uint16_t *)(buf + 6));
  header->size_y = ntohs(*(uint16_t *)(buf + 8));
  header->size_z = ntohs(*(uint16_t *)(buf + 10));

  /* Read slice probabilities */
  header->slice_probs = calloc_or_fail(header->size_y, 1);
  read_or_fail(fd, header->slice_probs, header->size_y);

  /* Read name/ID map */
  read_or_fail(fd, &header->name_id_count, 2);
  header->name_id_count = ntohs(header->name_id_count);
  /* First time to get the total length */
  offset = 2 * header->name_id_count;
  for (uint16_t i = 0; i < header->name_id_count; i++) {
    uint16_t name_length;
    read_or_fail(fd, &name_length, 2);
    name_length = ntohs(name_length);
    offset += name_length;
    lseek(fd, name_length, SEEK_CUR);
  }
  header->name_id_map = calloc_or_fail(offset, 1);
  lseek(fd, -offset, SEEK_CUR);
  /* Second time to actually read it */
  offset = 0;
  for (uint16_t i = 0; i < header->name_id_count; i++) {
    uint16_t name_length;
    read_or_fail(fd, &name_length, 2);
    name_length = ntohs(name_length);
    header->name_id_map[offset] = name_length;
    offset += 2;
    read_or_fail(fd, header->name_id_map + offset, name_length);
    offset += name_length;
  }
}

void read_nodedefs(int fd, mts *schematic) {
  uint8_t *source;
  void *dest;
  off_t start, end, source_length;
  size_t schematic_size, dest_length;

  /* Allocate memory for param0-2 */
  schematic_size =
    schematic->header.size_x * schematic->header.size_y * schematic->header.size_z;
  dest_length = 4 * schematic_size;
  dest = calloc_or_fail(dest_length, 1);
  schematic->param0 = dest;
  schematic->param1 = (uint8_t *)dest + 2 * schematic_size;
  schematic->param2 = (uint8_t *)dest + 3 * schematic_size;

  /* Calculate length of compressed data */
  start = lseek(fd, 0, SEEK_CUR);
  end = lseek(fd, 0, SEEK_END);
  source_length = end - start;
  lseek(fd, start, SEEK_SET);

  /* Read and uncompress compressed data */
  source = calloc_or_fail(source_length, 1);
  read_or_fail(fd, source, source_length);
  uncompress(dest, &dest_length, source, source_length);
  free(source);

  /* Convert param0 byte order */
  for (size_t i = 0; i < schematic_size; i++) {
    schematic->param0[i] = ntohs(schematic->param0[i]);
  }
}

void free_header(mts_header *header) {
  free(header->slice_probs);
  free(header->name_id_map);
}

void free_nodedefs(mts *schematic) {
  free(schematic->param0);
}

/* param1-2 indexing functions */
#define DEFINE_INDEX(type, n)                                           \
  type index_param##n(mts *schematic, uint16_t x, uint16_t y, uint16_t z) { \
    return schematic->param##n                                          \
      [z * schematic->header.size_z                                     \
       * schematic->header.size_y + y                                   \
       * schematic->header.size_x + x];                                 \
  }
DEFINE_INDEX(uint16_t, 0)
DEFINE_INDEX(uint8_t, 1)
DEFINE_INDEX(uint8_t, 2)
#undef DEFINE_INDEX

bool lookup_id_by_name(mts_header *header, uint16_t *id, char *name) {
  for (uint16_t i = 0, offset = 0; i < header->name_id_count; i++) {
    uint16_t name_length = (uint16_t) *(header->name_id_map + offset);
    offset += 2;
    if (strlen(name) == name_length &&
        !memcmp(name, (char *)(header->name_id_map + offset), name_length)) {
      *id = i;
      return true;
    }
    offset += name_length;
  }
  return false;
}

size_t lookup_name_by_id(mts_header *header, uint16_t id, char **name) {
  for (uint16_t i = 0, offset = 0; i < header->name_id_count; i++) {
    uint16_t name_length = (uint16_t) *(header->name_id_map + offset);
    offset += 2;
    if (i == id) {
      *name = (char *)(header->name_id_map + offset);
      return name_length;
    }
    offset += name_length;
  }
  /* Return "?" as name if ID is invalid */
  *name = "?";
  return 2;
}

void read_colour(char *s, colour *c) {
  uint32_t i = strtol(s, NULL, 16);
  while (i > 0xffffffff)
    i = i >> 4;
  /* Set alpha to 255 if unspecified */
  if (i <= 0xffffff)
    i = (i << 8) ^ 0x000000ff;
  c->red = (i & 0xff000000) >> 24;
  c->green = (i & 0x00ff0000) >> 16;
  c->blue = (i & 0x0000ff00) >> 8;
  c->alpha = i & 0x000000ff;
}

colour *read_colourmap(FILE *stream, mts_header *header) {
  char buf[BUFSIZE] = {0};
  colour *colourmap = calloc_or_fail(header->name_id_count, sizeof(colour));
  /* Defines whether or not an ID has had a colour specifically set,
     or if it's using the default, so changing the default doesn't
     override specific colours. */
  uint8_t *override = calloc_or_fail(header->name_id_count, sizeof(uint8_t));
  uint32_t line_number = 0;

  while (fgets(buf, BUFSIZE, stream)) {
    colour c;
    uint16_t id;
    char *ptr = buf;
    line_number++;

    /* Skip over leading whitespace */
    while (is_whitespace(*ptr)) ptr++;
    /* Skip over empty lines */
    if (is_newline(*ptr)) continue;

    /* Read name */
    if (!(ptr = strtok(ptr, " \t"))) continue;
    if (*ptr == '#') continue;  /* skip over comments */
    else if (!lookup_id_by_name(header, &id, ptr)) {
      /* '?' = wildcard (sets default colour) */
      if (strcmp(ptr, "?")) {
        fprintf(stderr, "Invalid name: '%s' on line %"PRIu32"\n", ptr, line_number);
        continue;
      }
      override[id] = 0;
    } else override[id] = 1;

    /* Read colour */
    if (!(ptr = strtok(NULL, " \t"))) continue;
    read_colour(ptr, &c);

    /* Set colour */
    if (override[id])
      colourmap[id] = c;
    else
      /* Set default colour */
      for (id = 0; id < header->name_id_count; id++)
        if (!override[id])
          colourmap[id] = c;
  }

  free(override);
  return colourmap;
}

void write_pnm(FILE *f, image *img, enum image_format imgfmt) {
  bool text = imgfmt == IMGFMT_PPM_TEXT;
  bool alpha = false;
  if (imgfmt == IMGFMT_PAM) {
    for (uint16_t i = 0; i < img->colours; i++) {
      if (img->colourmap[i].alpha != 255) {
        alpha = true;
        break;
      }
    }
    fprintf(f, "P7\nWIDTH %"PRIu16"\nHEIGHT %"PRIu16"\nDEPTH %d\n"
            "MAXVAL 255\nTUPLTYPE %s\nENDHDR\n",
            img->width, img->height,
            alpha ? 4 : 3, alpha ? "RGB_ALPHA" : "RGB");
  } else
    fprintf(f, "P%c\n%"PRIu16" %"PRIu16"\n255\n",
            text ? '3' : '6',
            img->width, img->height);
  for (uint16_t x = 0; x < img->width; x++) {
    for (uint16_t z = 0; z < img->height; z++) {
      colour c = img->colourmap[img->pixmap[img->width * x + z]];
      fprintf(f, text ? "%"PRIu8" %"PRIu8" %"PRIu8" " : "%c%c%c", c.red, c.green, c.blue);
      if (alpha) fprintf(f, "%c", c.alpha);
    }
    if (text) fprintf(f, "\n");
  }
}

void print_node(mts *schematic, uint16_t x, uint16_t y, uint16_t z) {
  char *name;
  uint16_t param0 = index_param0(schematic, x, y, z);
  uint16_t length = lookup_name_by_id(&schematic->header, param0, &name);
  printf("(%"PRIu16", %"PRIu16", %"PRIu16"): param0: %"PRIu16" (%.*s), param1: %"PRIu8", param2: %"PRIu8"\n",
         x, y, z,
         param0, length, name,
         index_param1(schematic, x, y, z),
         index_param2(schematic, x, y, z));
}

int cmd_header(int argc, char *argv[], mts *schematic) {
  uint32_t signature = htonl(schematic->header.signature);

  printf("HEADER:\n"
         "\tsignature = %"PRIx32" ('%.4s')\n"
         "\tversion = %"PRIu16"\n"
         "\tsize X = %"PRIu16"\n"
         "\tsize Y = %"PRIu16"\n"
         "\tsize Z = %"PRIu16"\n"
         "SLICE PROBABILITIES:\n",
         schematic->header.signature, (char *)&signature, schematic->header.version,
         schematic->header.size_x, schematic->header.size_y, schematic->header.size_z);

  for (uint16_t i = 0; i < schematic->header.size_y; i++) {
    printf("\t%"PRIu16": %"PRIu8" (%.2f%%)\n", i,
           schematic->header.slice_probs[i],
           (double)schematic->header.slice_probs[i] / 127 * 100);
  }

  printf("NAME/ID MAP (%"PRIu16" entries):\n", schematic->header.name_id_count);
  for (uint16_t i = 0, offset = 0; i < schematic->header.name_id_count; i++) {
    uint16_t name_length = (uint16_t) *(schematic->header.name_id_map + offset);
    offset += 2;
    printf("\t%"PRIu16": %.*s\n", i, name_length, schematic->header.name_id_map + offset);
    offset += name_length;
  }
  return 0;
}

int cmd_dumpall(int argc, char *argv[], mts *schematic) {
  foreach_node(schematic, x, y, z)
    print_node(schematic, x, y, z);
  return 0;
}

int cmd_dump(int argc, char *argv[], mts *schematic) {
  int x, y, z;
  x = atoi(argv[1]);
  y = atoi(argv[2]);
  z = atoi(argv[3]);
  if (x < 0 || x >= schematic->header.size_x ||
      y < 0 || y >= schematic->header.size_y ||
      z < 0 || z >= schematic->header.size_z) {
    fprintf(stderr, "Coordinates must be between (0, 0, 0) and (%"PRIu16", %"PRIu16", %"PRIu16")\n",
            schematic->header.size_x - 1,
            schematic->header.size_y - 1,
            schematic->header.size_z - 1);
    return 1;
  }
  print_node(schematic, x, y, z);
  return 0;
}

int cmd_list(int argc, char *argv[], mts *schematic) {
  uint16_t id;
  if (!lookup_id_by_name(&schematic->header, &id, argv[1])) {
    fprintf(stderr, "Invalid name: '%s'\n", argv[1]);
    return 1;
  }
  foreach_node(schematic, x, y, z)
    if (index_param0(schematic, x, y, z) == id)
      print_node(schematic, x, y, z);
  return 0;
}

int cmd_countall(int argc, char *argv[], mts *schematic) {
  uint64_t *count = calloc_or_fail(schematic->header.name_id_count, sizeof(uint64_t));
  foreach_node(schematic, x, y, z)
    count[index_param0(schematic, x, y, z)]++;
  for (uint16_t id = 0; id < schematic->header.name_id_count; id++) {
    char *name;
    uint16_t length = lookup_name_by_id(&schematic->header, id, &name);
    printf("%"PRIu64"\t%.*s\n", count[id], length, name);
  }
  free(count);
  return 0;
}

int cmd_count(int argc, char *argv[], mts *schematic) {
  uint16_t id;
  uint64_t count = 0;
  if (!lookup_id_by_name(&schematic->header, &id, argv[1])) {
    fprintf(stderr, "Invalid name: '%s'\n", argv[1]);
    return 1;
  }
  foreach_node(schematic, x, y, z)
    if (index_param0(schematic, x, y, z) == id)
      count++;
  printf("%"PRIu64"\n", count);
  return 0;
}

int cmd_map(int argc, char *argv[], mts *schematic) {
  image img;
  enum image_format imgfmt;
  int status = 0;
  int y = atoi(argv[1]);
  if (y < 0 || y >= schematic->header.size_y) {
    fprintf(stderr, "Depth must be between 0 and %"PRIu16"\n", schematic->header.size_y - 1);
    return 1;
  }

  /* Make image */
  img.width = schematic->header.size_x;
  img.height = schematic->header.size_z;
  img.colours = schematic->header.name_id_count;
  img.pixmap = calloc_or_fail(img.width * img.height, sizeof(uint16_t));
  img.colourmap = read_colourmap(stdin, &schematic->header);
  for (uint16_t x = 0; x < img.width; x++) {
    for (uint16_t z = 0; z < img.height; z++) {
      img.pixmap[img.width * x + z] = index_param0(schematic, x, y, z);
    }
  }

  /* Write image */
  if (argc > 2) {
    char *format = argv[2];
    if (!strcmp(format, "ppm-binary"))
      imgfmt = IMGFMT_PPM_BINARY;
    else if (!strcmp(format, "ppm-text"))
      imgfmt = IMGFMT_PPM_TEXT;
    else if (!strcmp(format, "pam"))
      imgfmt = IMGFMT_PAM;
    else {
      fprintf(stderr, "Unknown format: %s (expected ppm-binary, ppm-text, or pam)\n", argv[2]);
      status = 1;
      goto cleanup;
    }
  } else imgfmt = IMGFMT_PPM_TEXT;
  write_pnm(stdout, &img, imgfmt);
 cleanup:
  free(img.pixmap);
  free(img.colourmap);
  return status;
}

struct {
  const char *name;
  int min_argc;
  const char *params;
  int (*command)(int, char **, mts *);
  bool require_nodedefs;
  const char *usage;
} commands[] = {
  { "header", 0, NULL, cmd_header, false,
    "Print header (including name/ID map)." },
  { "dumpall", 0, NULL, cmd_dumpall, true,
    "Dump param0-2 for all nodes." },
  { "dump", 3, "<x> <y> <z>", cmd_dump, true,
    "Dump param0-2 for a single node." },
  { "list", 1, "<name>", cmd_list, true,
    "List all nodes of a certain type by name." },
  { "countall", 0, NULL, cmd_countall, true,
    "Count all nodes by name." },
  { "count", 1, "<name>", cmd_count, true,
    "Count all nodes of a certain type by name." },
  { "map", 1, "<depth> [<format>]", cmd_map, true,
    "Make an image map of all nodes in a single layer.\n"
    "Takes name->colour map from stdin.\n"
    "<format> may be ppm-text, ppm-binary, or pam. Defaults to ppm-text." },
};

void usage(int status) {
  fprintf(stderr,
          "Usage: mtsread <file> <command> [args]\n"
          "Commands:\n");
  for (size_t i = 0; i < lengthof(commands); i++) {
    fprintf(stderr, "  %s", commands[i].name);
    if (commands[i].params)
      fprintf(stderr, " %s", commands[i].params);
    fprintf(stderr, "\n    ");
    for (size_t j = 0; commands[i].usage[j]; j++) {
      char c = commands[i].usage[j];
      fprintf(stderr, "%c", c);
      if (is_newline(c))
        fprintf(stderr, "    ");
    }
    fprintf(stderr, "\n");
  }
  exit(status);
}

int main(int argc, char *argv[]) {
  int fd;
  mts schematic;

  if (argc < 3) usage(1);
  if ((fd = open(argv[1], O_RDONLY)) == -1) {
    perror("open");
    return 1;
  }

  read_header(fd, &schematic.header);

  /* Check signature and version */
  if (schematic.header.signature != MTSCHEM_FILE_SIGNATURE) {
    fprintf(stderr,
            "Error: Invalid signature %"PRIx32" (expected %"PRIx32"). "
            "Are you sure this is a Minetest schematic file?\n",
            schematic.header.signature, MTSCHEM_FILE_SIGNATURE);
    return 1;
  }
  if (schematic.header.version != MTSCHEM_FILE_VERSION)
    fprintf(stderr,
            "Warning: Unsupported version number %"PRIu16" (expected %"PRIu16").\n",
            schematic.header.version, MTSCHEM_FILE_VERSION);

  for (size_t i = 0; i < lengthof(commands); i++) {
    int cargc = argc - 2;
    char **cargv = &argv[2];
    if (!strcmp(cargv[0], commands[i].name) && cargc > commands[i].min_argc) {
      int status;
      bool require_nodedefs = commands[i].require_nodedefs;
      if (require_nodedefs)
        read_nodedefs(fd, &schematic);
      status = commands[i].command(cargc, cargv, &schematic);
      if (require_nodedefs)
        free_nodedefs(&schematic);
      free_header(&schematic.header);
      return status;
    }
  }
  free_header(&schematic.header);
  usage(1);
  /* Should be unreachable */
  return 1;
}