Embed the data files needed for the ssh kitten into the Go binary

gen-go-code.py

@@ -641,6 +641,11 @@ def generate_textual_mimetypes() -> str:
     return '\n'.join(ans)
 
 
+def write_compressed_data(data: bytes, d: BinaryIO) -> None:
+    d.write(struct.pack('<I', len(data)))
+    d.write(zlib.compress(data, zlib.Z_BEST_COMPRESSION))
+
+
 def generate_unicode_names(src: TextIO, dest: BinaryIO) -> None:
     num_names, num_of_words = map(int, next(src).split())
     gob = io.BytesIO()
@@ -655,9 +660,32 @@ def generate_unicode_names(src: TextIO, dest: BinaryIO) -> None:
             if aliases:
                 record += aliases.encode()
             gob.write(struct.pack('<H', len(record)) + record)
-    data = gob.getvalue()
-    dest.write(struct.pack('<I', len(data)))
-    dest.write(zlib.compress(data, zlib.Z_BEST_COMPRESSION))
+    write_compressed_data(gob.getvalue(), dest)
+
+
+def generate_ssh_kitten_data() -> None:
+    files = {
+        'terminfo/kitty.terminfo', 'terminfo/x/xterm-kitty',
+    }
+    for dirpath, dirnames, filenames in os.walk('shell-integration'):
+        for f in filenames:
+            path = os.path.join(dirpath, f)
+            files.add(path.replace(os.sep, '/'))
+    dest = 'tools/cmd/ssh/data_generated.bin'
+    if newer(dest, *files):
+        buf = io.BytesIO()
+        fmap = dict.fromkeys(files, (0, 0))
+        for f in fmap:
+            with open(f, 'rb') as src:
+                data = src.read()
+            pos = buf.tell()
+            buf.write(data)
+            size = len(data)
+            fmap[f] = pos, size
+        mapping = ','.join(f'{name} {pos[0]} {pos[1]}' for name, pos in fmap.items()).encode('ascii')
+        data = struct.pack('<I', len(fmap)) + mapping + b'\n' + buf.getvalue()
+        with open(dest, 'wb') as d:
+            write_compressed_data(data, d)
 
 
 def main() -> None:
@@ -676,6 +704,7 @@ def main() -> None:
     if newer('tools/unicode_names/data_generated.bin', 'tools/unicode_names/names.txt'):
         with open('tools/unicode_names/data_generated.bin', 'wb') as dest, open('tools/unicode_names/names.txt') as src:
             generate_unicode_names(src, dest)
+    generate_ssh_kitten_data()
 
     update_completion()
     update_at_commands()

tools/cmd/ssh/data.go

@@ -0,0 +1,37 @@
+// License: GPLv3 Copyright: 2023, Kovid Goyal, <kovid at kovidgoyal.net>
+
+package ssh
+
+import (
+	"bytes"
+	_ "embed"
+	"encoding/binary"
+	"fmt"
+	"kitty/tools/utils"
+	"strconv"
+	"strings"
+)
+
+var _ = fmt.Print
+
+//go:embed data_generated.bin
+var embedded_data string
+
+type Container = map[string][]byte
+
+var Data = (&utils.Once[Container]{Run: func() Container {
+	raw := utils.ReadCompressedEmbeddedData(embedded_data)
+	num_of_entries := binary.LittleEndian.Uint32(raw)
+	raw = raw[4:]
+	ans := make(Container, num_of_entries)
+	idx := bytes.IndexByte(raw, '\n')
+	text := utils.UnsafeBytesToString(raw[:idx])
+	raw = raw[idx+1:]
+	for _, record := range strings.Split(text, ",") {
+		parts := strings.Split(record, " ")
+		offset, _ := strconv.Atoi(parts[1])
+		size, _ := strconv.Atoi(parts[2])
+		ans[parts[0]] = raw[offset : offset+size]
+	}
+	return ans
+}}).Get

tools/cmd/ssh/main.go

@@ -170,6 +170,7 @@ type connection_data struct {
 }
 
 func run_ssh(ssh_args, server_args, found_extra_args []string) (rc int, err error) {
+	go Data()
 	cmd := append([]string{SSHExe()}, ssh_args...)
 	cd := connection_data{remote_args: server_args[1:]}
 	hostname := server_args[0]

tools/unicode_names/query.go

@@ -4,11 +4,9 @@ package unicode_names
 
 import (
 	"bytes"
-	"compress/zlib"
 	_ "embed"
 	"encoding/binary"
 	"fmt"
-	"io"
 	"strings"
 	"sync"
 	"time"
@@ -64,33 +62,8 @@ func parse_record(record []byte, mark uint16) {
 
 var parse_once sync.Once
 
-func read_all(r io.Reader, expected_size int) ([]byte, error) {
-	b := make([]byte, 0, expected_size)
-	for {
-		if len(b) == cap(b) {
-			// Add more capacity (let append pick how much).
-			b = append(b, 0)[:len(b)]
-		}
-		n, err := r.Read(b[len(b):cap(b)])
-		b = b[:len(b)+n]
-		if err != nil {
-			if err == io.EOF {
-				err = nil
-			}
-			return b, err
-		}
-	}
-}
-
 func parse_data() {
-	compressed := utils.UnsafeStringToBytes(unicode_name_data)
-	uncompressed_size := binary.LittleEndian.Uint32(compressed)
-	r, _ := zlib.NewReader(bytes.NewReader(compressed[4:]))
-	defer r.Close()
-	raw, err := read_all(r, int(uncompressed_size))
-	if err != nil {
-		panic(err)
-	}
+	raw := utils.ReadCompressedEmbeddedData(unicode_name_data)
 	num_of_lines := binary.LittleEndian.Uint32(raw)
 	raw = raw[4:]
 	num_of_words := binary.LittleEndian.Uint32(raw)

tools/utils/embed.go

@@ -0,0 +1,43 @@
+// License: GPLv3 Copyright: 2023, Kovid Goyal, <kovid at kovidgoyal.net>
+
+package utils
+
+import (
+	"bytes"
+	"compress/zlib"
+	"encoding/binary"
+	"fmt"
+	"io"
+)
+
+var _ = fmt.Print
+
+func ReadAll(r io.Reader, expected_size int) ([]byte, error) {
+	b := make([]byte, 0, expected_size)
+	for {
+		if len(b) == cap(b) {
+			// Add more capacity (let append pick how much).
+			b = append(b, 0)[:len(b)]
+		}
+		n, err := r.Read(b[len(b):cap(b)])
+		b = b[:len(b)+n]
+		if err != nil {
+			if err == io.EOF {
+				err = nil
+			}
+			return b, err
+		}
+	}
+}
+
+func ReadCompressedEmbeddedData(raw string) []byte {
+	compressed := UnsafeStringToBytes(raw)
+	uncompressed_size := binary.LittleEndian.Uint32(compressed)
+	r, _ := zlib.NewReader(bytes.NewReader(compressed[4:]))
+	defer r.Close()
+	ans, err := ReadAll(r, int(uncompressed_size))
+	if err != nil {
+		panic(err)
+	}
+	return ans
+}