minetest-goutils/schematic/schematic.go

// Copyright 2023 prestidigitator (as registered on forum.minetest.net)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package schematic

import (
    "bytes"
    "compress/zlib"
    "encoding/binary"
    "encoding/json"
    "fmt"
    "io"
    "strings"
    "text/template"

    "gopkg.in/yaml.v3"
)


var yamlTemplate = template.Must(template.New("yaml").Funcs(map[string]any{
    "yProbVal":    func(prob   uint8) uint8 { return (prob   & 0x7f) << 1 },
    "param1Prob":  func(param1 uint8) uint8 { return (param1 & 0x7f) << 1 },
    "param1Force": func(param1 uint8) uint8 { return  param1 & 0x80 },
    "strEsc": func(s string) (string, error) {
        bs, err := yaml.Marshal(s); return strings.TrimSpace(string(bs)), err
    },
}).Parse(strings.TrimSpace(`
    size: {x: {{ .Size.X }}, y: {{ .Size.Y }}, z: {{ .Size.Z -}} }

    {{- "\n" -}}

    yslice_prob:
        {{- range $y, $p := .YProbs -}}
            {{- "\n" -}}
            - {ypos: {{ $y }}, prob: {{ yProbVal $p -}} }
        {{- end -}}

    {{- "\n" -}}

    data:
        {{- range $_, $node := .Nodes -}}
            {{- $n  := index $.Names $node.Content | strEsc -}}
            {{- $p  := param1Prob  $node.Param1 -}}
            {{- $f  := param1Force $node.Param1 -}}
            {{- $p2 := $node.Param2 -}}

            {{- "\n" -}}
            - {name: {{ $n }}, prob: {{ $p }}, param2: {{ $p2 }}
                 {{- if $f -}} , force_place: true {{- end -}}
              }
        {{- end -}}

    {{- "\n" -}}
`)))

var jsonTemplate = template.Must(template.New("json").Funcs(map[string]any{
    "yProbVal":    func(prob   uint8) uint8 { return (prob   & 0x7f) << 1 },
    "param1Prob":  func(param1 uint8) uint8 { return (param1 & 0x7f) << 1 },
    "param1Force": func(param1 uint8) uint8 { return  param1 & 0x80 },
    "strEsc": func(s string) (string, error) {
        bs, err := json.Marshal(s); return string(bs), err
    },
}).Parse(strings.TrimSpace(`
    {{- $s      := .schematic -}}
    {{- $indent := .indent    -}}

    {{- $i0   := ""  -}}
    {{- $i1oc := ""  -}}
    {{- $i1   := " " -}}
    {{- $i2oc := ""  -}}
    {{- $i2   := " " -}}
    {{- if $indent -}}
        {{- $i0   =        "\n"                     -}}
        {{- $i1oc = printf "\n%s"   $indent         -}}
        {{- $i1   = printf "\n%s"   $indent         -}}
        {{- $i2oc = printf "\n%s%s" $indent $indent -}}
        {{- $i2   = printf "\n%s%s" $indent $indent -}}
    {{- end -}}

    {
        {{- $i1oc -}}

        "size": {"x": {{ $s.Size.X }}, "y": {{ $s.Size.Y }}, "z": {{ $s.Size.Z -}} },

        {{- $i1 -}}

        "yslice_prob": [
            {{- range $y, $p := $s.YProbs -}}
                {{- if gt $y 0 -}} ,{{- $i2 -}} {{- else -}}{{- $i2oc -}}{{- end -}}
                {"ypos": {{ $y }}, "prob": {{ yProbVal $p -}} }
            {{- end -}}

            {{- $i1oc -}}
        ],

        {{- $i1 -}}

        "data": [

            {{- range $i, $node := $s.Nodes -}}
                {{- $n  := index $s.Names $node.Content | strEsc -}}
                {{- $p  := param1Prob  $node.Param1 -}}
                {{- $f  := param1Force $node.Param1 -}}
                {{- $p2 := $node.Param2 -}}

                {{- if gt $i 0 -}} ,{{- $i2 -}} {{- else -}}{{- $i2oc -}}{{- end -}}
                {"name": {{ $n }}, "prob": {{ $p }}, "param2": {{ $p2 }}
                    {{- if $f -}} , "force_place": true {{- end -}}
                }
            {{- end -}}

            {{- $i1oc -}}
        ]

        {{- $i0 -}}
    }
    {{- $i0 -}}
`)))


var luaTemplate = template.Must(template.New("lua").Funcs(map[string]any{
    "inc":         func(n      int)   int   { return n+1 },
    "yProbVal":    func(prob   uint8) uint8 { return (prob   & 0x7f) << 1 },
    "param1Prob":  func(param1 uint8) uint8 { return (param1 & 0x7f) << 1 },
    "param1Force": func(param1 uint8) uint8 { return  param1 & 0x80 },
    "strEsc": func(s string) (string, error) {
        var sb strings.Builder
        sb.Grow(len(s)+2)
        sb.WriteRune('"')
        for i := 0; i < len(s); i++ {
            switch b := s[i]; b {
                case '\a': sb.WriteString(`\a`)
                case '\b': sb.WriteString(`\b`)
                case '\f': sb.WriteString(`\f`)
                case '\n': sb.WriteString(`\n`)
                case '\r': sb.WriteString(`\r`)
                case '\t': sb.WriteString(`\t`)
                case '\\': sb.WriteString(`\\`)
                case '"':  sb.WriteString(`\"`)
                case '\'': sb.WriteString(`\'`)
                case 0:    sb.WriteString(`\0`)
                default:
                    if 0x20 <= b && b <= 0x7e {
                        sb.WriteByte(b)
                    } else {
                        fmt.Fprintf(&sb, `\%03d`, b)
                    }
            }
        }
        sb.WriteRune('"')
        return sb.String(), nil
    },
}).Parse(strings.TrimSpace(`
    {{- $s        := .schematic -}}
    {{- $indent   := .indent    -}}
    {{- $comments := .comments  -}}

    schematic = {
        {{- printf "\n%s" $indent -}}

        size = {x={{ printf "%d" $s.Size.X }}, y={{ printf "%d" $s.Size.Y -}}
              , z={{ printf "%d" $s.Size.Z -}} },

        {{- printf "\n%s" $indent -}}

        yslice_prob = {
            {{- range $y, $p := $s.YProbs -}}
                {{- printf "\n%s%s" $indent $indent -}}
                {ypos={{ printf "%d" $y }}, prob={{ yProbVal $p | printf "%d" -}} },
            {{- end -}}

            {{- printf "\n%s" $indent -}}
        },

        {{- printf "\n%s" $indent -}}

        data = {
            {{- $x := 0 -}}
            {{- $y := 0 -}}
            {{- $z := 0 -}}
            {{- range $i, $node := $s.Nodes}}
                {{- $x = inc $x -}}
                {{- if ge $x $s.Size.X -}}
                    {{- $x = 0 -}}
                    {{- $y = inc $y -}}
                    {{- if ge $y $s.Size.Y -}}
                        {{- $y = 0 -}}
                        {{- $z = inc $z -}}
                    {{- end -}}
                {{- end -}}

                {{- $n  := index $s.Names $node.Content | strEsc -}}
                {{- $p  := param1Prob  $node.Param1 -}}
                {{- $f  := param1Force $node.Param1 -}}
                {{- $p2 := $node.Param2 -}}

                {{- if and $comments (eq $x 0) -}}
                    {{- if or (gt $y 0) (gt $z 0) -}}
                        {{- "\n" -}}
                    {{- end -}}
                    {{- printf "\n%s%s" $indent $indent -}}
                    -- z={{ $z }}, y={{ $y }}
                {{- end -}}

                {{- printf "\n%s%s" $indent $indent -}}
                {name={{ $n }}, prob={{ printf "%d" $p }}, param2={{ printf "%d" $p2 }}
                    {{- if $f -}} , force_place=true {{- end -}}
                },
            {{- end -}}

            {{- printf "\n%s" $indent -}}
        }
        {{- "\n" -}}
    }
    {{- "\n" -}}
`)))


type Schematic struct{
    Size   Vec3i
    YProbs []uint8
    Names  []string
    Nodes  []Node
}

type Vec3i struct{
    X, Y, Z uint16
}

type Node struct{
    Content        uint16
    Param1, Param2 uint8
}

type schematicSer struct{
    Size   sizeSer    `yaml:"size"        json:"size"`
    YProbs []yprobSer `yaml:"yslice_prob" json:"yslice_prob"`
    Data   []datumSer `yaml:"data"        json:"data"`
}

type sizeSer struct{
    X uint16 `yaml:"x" json:"x"`
    Y uint16 `yaml:"y" json:"y"`
    Z uint16 `yaml:"z" json:"z"`
}

type yprobSer struct{
    YPos uint16 `yaml:"ypos" json:"ypos"`
    Prob uint8  `yaml:"prob" json:"prob"`
}

type datumSer struct{
    Name   string `yaml:"name"                  json:"name"`
    Prob   uint8  `yaml:"prob"                  json:"prob"`
    Param2 uint8  `yaml:"param2"                json:"param2"`
    Force  bool   `yaml:"force_place,omitempty" json:"force_place,omitempty"`
}

func FromMTSStream(ins io.Reader) (s Schematic, err error) {
    defer func() {
        if pval := recover(); pval != nil && pval != &err { panic(pval) }
    }()

    buf := make([]byte, 12)
    readnb := func(ins io.Reader, n int) {
        if len(buf) < n { buf = make([]byte, n) }
        _, err = io.ReadFull(ins, buf[:n])
        if err != nil { panic(&err) }
    }

    readnb(ins, 12)
    sig := append([]byte{}, buf[:4]...)
    ver := binary.BigEndian.Uint16(buf[4:6])
    sx  := binary.BigEndian.Uint16(buf[6:8])
    sy  := binary.BigEndian.Uint16(buf[8:10])
    sz  := binary.BigEndian.Uint16(buf[10:12])
    if bytes.Compare(sig, []byte("MTSM")) != 0 { err = InvalidMTSSignature;   return }
    if ver > 4                                 { err = UnsupportedMTSVersion; return }

    readnb(ins, int(sy))
    yprobs := make([]uint8, sy)
    if ver >= 3 {
        for i, _ := range yprobs { yprobs[i] = buf[i] }
    } else {
        for i, _ := range yprobs { yprobs[i] = 255 }
    }

    readnb(ins, 2)
    nnames := binary.BigEndian.Uint16(buf[:2])

    names, ignoreIndex := make([]string, nnames), -1
    for i, _ := range names {
        readnb(ins, 2)
        nameLen := int(binary.BigEndian.Uint16(buf[:2]))

        readnb(ins, nameLen)
        name := string(buf[:nameLen])

        if name == "ignore" { name, ignoreIndex = "air", i }
        names[i] = name
    }

    vol := sx * sy * sz

    zin, err := zlib.NewReader(ins)
    if err != nil { return }
    defer zin.Close()

    readnb(zin, 2*int(vol))
    content := make([]uint16, vol)
    for i, _ := range content {
        content[i] = binary.BigEndian.Uint16(buf[2*i : 2*i+2])
    }

    readnb(zin, int(vol))
    param1s := make([]uint8, vol)
    for i, _ := range param1s { param1s[i] = buf[i] }

    readnb(zin, int(vol))
    param2s := make([]uint8, vol)
    for i, _ := range param2s { param2s[i] = buf[i] }

    if ver < 2 {
        for i, _ := range param1s {
            if ignoreIndex >= 0 && int(content[i]) == ignoreIndex {
                param1s[i] = 0
            } else if param1s[i] == 0 {
                param1s[i] = 255
            }
        }
    }

    if ver < 4 {
        for i, _ := range yprobs  { yprobs[i]  >>= 1 }
        for i, _ := range param1s { param1s[i] >>= 1 }
    }

    nodes := make([]Node, vol)
    for i, _ := range nodes {
        nodes[i] = Node{
            Content: content[i],
            Param1:  param1s[i],
            Param2:  param2s[i],
        }
    }

    s = Schematic{
        Size:   Vec3i{X: sx, Y: sy, Z: sz},
        YProbs: yprobs,
        Names:  names,
        Nodes:  nodes,
    }
    return
}

func FromYAMLStream(ins io.Reader) (Schematic, error) {
    var ss schematicSer
    err := yaml.NewDecoder(ins).Decode(&ss)
    if err != nil { return Schematic{}, err }
    return fromSer(ss)
}

func FromJSONStream(ins io.Reader) (Schematic, error) {
    var ss schematicSer
    err := json.NewDecoder(ins).Decode(&ss)
    if err != nil { return Schematic{}, err }
    return fromSer(ss)
}

func fromSer(ss schematicSer) (s Schematic, err error) {
    sx, sy, sz := ss.Size.X, ss.Size.Y, ss.Size.Z
    s.Size.X, s.Size.Y, s.Size.Z = sx, sy, sz

    yprobs := make([]uint8, sy)
    for _, yp := range ss.YProbs {
        y, p := yp.YPos, yp.Prob
        if y >= sy {
            err = fmt.Errorf("y-slice y index %d out-of-range [0, %d)", y, sy)
            return
        }
        yprobs[y] = p >> 1
    }
    s.YProbs = yprobs

    names, nameMap := []string{}, map[string]uint16{}
    for _, d := range ss.Data {
        name := d.Name
        if _, ok := nameMap[name]; !ok {
            nameMap[name] = uint16(len(names))
            names         = append(names, name)
        }
    }
    s.Names = names

    nodes := make([]Node, len(ss.Data))
    for i, d := range ss.Data {
        param1 := d.Prob >> 1
        if d.Force { param1 |= 0x80 }

        nodes[i] = Node{
            Content: nameMap[d.Name],
            Param1:  param1,
            Param2:  d.Param2,
        }
    }
    s.Nodes = nodes

    return
}

func (s Schematic) WriteMTS(outs io.Writer) (err error) {
    var header    [12]byte
    var uint16Buf [2]byte

    copy(header[:6], []byte{'M', 'T', 'S', 'M', 0, 4})
    binary.BigEndian.PutUint16(header[6:8],   s.Size.X)
    binary.BigEndian.PutUint16(header[8:10],  s.Size.Y)
    binary.BigEndian.PutUint16(header[10:12], s.Size.Z)
    _, err = outs.Write(header[:])
    if err != nil { return err }

    for _, p := range s.YProbs {
        _, err = outs.Write([]byte{p})
        if err != nil { return err }
    }

    binary.BigEndian.PutUint16(uint16Buf[:2], uint16(len(s.Names)))
    _, err = outs.Write(uint16Buf[:])
    if err != nil { return err }

    for _, name := range s.Names {
        binary.BigEndian.PutUint16(uint16Buf[:2], uint16(len(name)))
        _, err = outs.Write(uint16Buf[:])
        if err != nil { return err }
        _, err = outs.Write([]byte(name))
        if err != nil { return err }
    }

    zout := zlib.NewWriter(outs)
    defer func() {
        if e := zout.Close(); e != nil && err == nil { err = e }
    }()

    for _, node := range s.Nodes {
        binary.BigEndian.PutUint16(uint16Buf[:2], uint16(node.Content))
        _, err = zout.Write(uint16Buf[:])
        if err != nil { return err }
    }
    for _, node := range s.Nodes {
        _, err = zout.Write([]byte{byte(node.Param1)})
        if err != nil { return err }
    }
    for _, node := range s.Nodes {
        _, err = zout.Write([]byte{byte(node.Param2)})
        if err != nil { return err }
    }
    return
}

func (s Schematic) WriteYAML(outs io.Writer, indent string) error {
    return yamlTemplate.Execute(outs, s)
}

func (s Schematic) WriteJSON(outs io.Writer, indent string) error {
    return jsonTemplate.Execute(outs, map[string]any{"indent": indent, "schematic": s})
}

func (s Schematic) WriteLua(outs io.Writer, indent string, comments bool) (
    err error,
) {
    return luaTemplate.Execute(
        outs, map[string]any{"schematic": s, "indent": indent, "comments": comments},
    )
}