Fetches terrain chunks from a LevelDB.

terrainfetch.go

//
// terrainfetch.go
// ---------------
//
// Fetches terrain chunks from a LevelDB.
//
// (c) 2012 by Sascha L. Teichmann
//
package main

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"flag"
	"fmt"
	leveldb "github.com/jmhodges/levigo"
	"io"
	"log"
	"os"
	//leveldb "github.com/jmhodges/levigo_leveldb_1.4"
)

const (
	BITS      = uint32(20)
	MSB       = 3*BITS - 1
	ZERO32    = uint32(0)
	ONE32     = uint32(1)
	ZERO64    = uint64(0)
	ONE64     = uint64(1)
	HI_MASK   = ONE32 << (BITS + ONE32)
	WORLD_MID = ^(^0 << 19)
	_000_     = 0
	_001_     = 1
	_010_     = 1 << 1
	_011_     = (1 << 1) | 1
	_100_     = 1 << 2
	_101_     = (1 << 2) | 1
	// MASK = hex(int(''.join(['1' if (i%3) == 0 else '0' for i in range(60)]),2))
	MASK uint64 = 0x924924924924924
	// FULL = hex(int('1'*60,2))
	FULL uint64 = 0xfffffffffffffff
)

func ZEncode(x, y, z uint32) uint64 {
	code, p := ZERO64, ONE64
	for mask := ONE32; mask != HI_MASK; mask <<= 1 {
		if (x & mask) != 0 {
			code |= p
		}
		p <<= 1
		if (y & mask) != 0 {
			code |= p
		}
		p <<= 1
		if (z & mask) != 0 {
			code |= p
		}
		p <<= 1
	}
	return code
}

func zencode(x, y, z int) uint64 {
	return ZEncode(
		uint32(x+WORLD_MID),
		uint32(y+WORLD_MID),
		uint32(z+WORLD_MID))
}

func ZCodeAsKey(code uint64) []byte {
	var b bytes.Buffer // XXX: Is there a better way todo this?
	binary.Write(&b, binary.BigEndian, &code)
	return b.Bytes()
}

func KeyAsZCode(key []byte) uint64 {
	b := bytes.NewBuffer(key)
	var code uint64
	binary.Read(b, binary.BigEndian, &code)
	return code
}

func ZDecode(code uint64) (x, y, z uint32) {
	p := ONE64
	x, y, z = ZERO32, ZERO32, ZERO32
	for mask := ONE32; mask != HI_MASK; mask <<= 1 {
		if (code & p) != 0 {
			x |= mask
		}
		p <<= 1
		if (code & p) != 0 {
			y |= mask
		}
		p <<= 1
		if (code & p) != 0 {
			z |= mask
		}
		p <<= 1
	}
	return
}

type Cube struct {
	X1, Y1, Z1 int
	X2, Y2, Z2 int
}

func (c *Cube) Order() {
	if c.X1 > c.X2 {
		c.X1, c.X2 = c.X2, c.X1
	}
	if c.Y1 > c.Y2 {
		c.Y1, c.Y2 = c.Y2, c.Y1
	}
	if c.Z1 > c.Z2 {
		c.Z1, c.Z2 = c.Z2, c.Z1
	}
}

func WorldCoords(zcode uint64) (int, int, int) {
	xc, yc, zc := ZDecode(zcode)
	return int(xc) - WORLD_MID, int(yc) - WORLD_MID, int(zc) - WORLD_MID
}

func (c *Cube) Contains(zcode uint64) bool {
	x, y, z := WorldCoords(zcode)
	return c.X1 <= x && x <= c.X2 &&
		c.Y1 <= y && y <= c.Y2 &&
		c.Z1 <= z && z <= c.Z2
}

func (c *Cube) ZEncoderP1() uint64 {
	return zencode(c.X1, c.Y1, c.Z1)
}

func (c *Cube) ZEncoderP2() uint64 {
	return zencode(c.X2, c.Y2, c.Z2)
}

type RLEDecoder struct {
	buf       []byte
	current   byte
	pos, left int
}

func NewRLEDecoder(buf []byte) *RLEDecoder {
	return &RLEDecoder{buf, 0, 0, 0}
}

func (r *RLEDecoder) ReadByte() (byte, error) {
	if r.left > 0 {
		r.left--
		return r.current, nil
	}
	if r.pos >= len(r.buf) {
		return 0, io.EOF
	}
	x := r.buf[r.pos]
	r.pos++
	if (x & (1 << 7)) == 0 {
		return x, nil
	}
	count := int(x & 127)
	if count == 0 {
		return 0, io.ErrUnexpectedEOF
	}
	r.left = count - 1
	if r.pos >= len(r.buf) {
		return 0, io.EOF
	}
	r.current = r.buf[r.pos]
	r.pos++
	return r.current, nil
}

func DumpData(zcode uint64, data []byte) error {
	x, y, z := WorldCoords(zcode)
	filename := fmt.Sprintf("tile_%d_%d_%d.raw", x, y, z)
	log.Printf("Writing: %s (%d)\n", filename, len(data))
	f, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer f.Close()
	rle := NewRLEDecoder(data)
	b := bufio.NewWriter(f)
	for {
		x, err := rle.ReadByte()
		if err != nil {
			if err == io.EOF {
				break
			}
			return err
		}
		b.WriteByte(x)
	}
	return b.Flush()
}

func setbits(p uint32, v uint64) uint64 {
	mask := (MASK >> (MSB - p)) & (^(FULL << p) & FULL)
	return (v | mask) & ^(1 << p) & FULL
}

func unsetbits(p uint32, v uint64) uint64 {
	mask := ^(MASK >> (MSB - p)) & FULL
	return (v & mask) | (1 << p)
}

func BigMin(minz, maxz, zcode uint64) uint64 {
	bigmin := maxz
	pos := MSB
	for mask := ONE64 << MSB; mask != 0; mask >>= 1 {
		var v int
		if (zcode & mask) != 0 {
			v = _100_
		} else {
			v = _000_
		}
		if (minz & mask) != 0 {
			v |= _010_
		}
		if (maxz & mask) != 0 {
			v |= _001_
		}

		switch v {
		case _001_:
			bigmin = unsetbits(pos, minz)
			maxz = setbits(pos, maxz)
		case _011_:
			return minz
		case _100_:
			return bigmin
		case _101_:
			minz = unsetbits(pos, minz)
		}
		pos--
	}
	return bigmin
}

func main() {
	var (
		db_path string
		query   Cube
	)

	flag.StringVar(&db_path, "db", "terrain.db", "World database path")
	flag.IntVar(&query.X1, "x1", -1, "x1 of query cube")
	flag.IntVar(&query.Y1, "y1", -1, "y1 of query cube")
	flag.IntVar(&query.Z1, "z1", -1, "z1 of query cube")
	flag.IntVar(&query.X2, "x2", 1, "x2 of query cube")
	flag.IntVar(&query.Y2, "y2", 1, "y2 of query cube")
	flag.IntVar(&query.Z2, "z2", 1, "z2 of query cube")

	flag.Parse()

	opts := leveldb.NewOptions()
	db, err := leveldb.Open(db_path, opts)
	if err != nil {
		log.Fatal("Cannot open database", err)
	}
	defer db.Close()

	ro := leveldb.NewReadOptions()

	query.Order()
	zmin, zmax := query.ZEncoderP1(), query.ZEncoderP2()

	it := db.NewIterator(ro)
	defer it.Close()

	it.Seek(ZCodeAsKey(zmin))

	for it.Valid() {
		zcode := KeyAsZCode(it.Key())
		//fmt.Printf("key: %x\n", zcode)
		if zcode > zmax {
			break
		}
		if query.Contains(zcode) {
			if err := DumpData(zcode, it.Value()); err != nil {
				log.Fatal("Cannot write data", err)
			}
			it.Next()
		} else {
			it.Seek(ZCodeAsKey(BigMin(zmin, zmax, zcode)))
		}
	}

	if err := it.GetError(); err != nil {
		log.Fatal("Seeking iterator failed", err)
	}
}