From a312577297526cfd3a51f68d8c112229f3c4185a Mon Sep 17 00:00:00 2001
From: jet2tlf <jet2tlf@gmail.com>
Date: Mon, 3 Jun 2024 14:29:00 -0300
Subject: codecrafters submit [skip ci]

---
 cmd/mybittorrent/decoder.go | 78 ++++++++++++++++++++++++++++++++++-----------
 1 file changed, 60 insertions(+), 18 deletions(-)

(limited to 'cmd/mybittorrent/decoder.go')

diff --git a/cmd/mybittorrent/decoder.go b/cmd/mybittorrent/decoder.go
index 10066e1..641b042 100644
--- a/cmd/mybittorrent/decoder.go
+++ b/cmd/mybittorrent/decoder.go
@@ -12,11 +12,11 @@ type Decoder struct {
 	decoderFuncs []DecoderFunc
 }
 
-type DecoderFunc func(string) (any, int, bool)
+type DecoderFunc func(string) (any, int, error)
 
 func NewDecoder(raw string) *Decoder {
 	d := &Decoder{raw: raw}
-	return d.WithDecoderFunc(d.decodeSlice).WithDecoderFunc(d.decodeString).WithDecoderFunc(d.decodeInteger)
+	return d.WithDecoderFunc(d.decodeSlice).WithDecoderFunc(d.decodeMap).WithDecoderFunc(d.decodeString).WithDecoderFunc(d.decodeInteger)
 }
 
 func (d *Decoder) WithDecoderFunc(f DecoderFunc) *Decoder {
@@ -26,7 +26,7 @@ func (d *Decoder) WithDecoderFunc(f DecoderFunc) *Decoder {
 
 func (d *Decoder) Decode() (any, error) {
 	for _, f := range d.decoderFuncs {
-		if r, _, ok := f(d.raw); ok {
+		if r, l, err := f(d.raw); l > 0 || err != nil {
 			return r, nil
 		}
 	}
@@ -34,9 +34,9 @@ func (d *Decoder) Decode() (any, error) {
 	return nil, fmt.Errorf("failed to decode value: %s", d.raw)
 }
 
-func (d *Decoder) decodeString(chunk string) (any, int, bool) {
+func (d *Decoder) decodeString(chunk string) (any, int, error) {
 	if !unicode.IsDigit(rune(chunk[0])) {
-		return "", 0, false
+		return "", 0, nil
 	}
 
 	var firstColonIndex int
@@ -52,47 +52,89 @@ func (d *Decoder) decodeString(chunk string) (any, int, bool) {
 
 	length, err := strconv.Atoi(lengthStr)
 	if err != nil {
-		return "", 0, false
+		return "", 0, err
 	}
 
 	val := chunk[firstColonIndex+1 : firstColonIndex+1+length]
 
-	return val, length + 1 + len(lengthStr), true
+	return val, length + 1 + len(lengthStr), nil
 }
 
-func (d *Decoder) decodeInteger(chunk string) (any, int, bool) {
+func (d *Decoder) decodeInteger(chunk string) (any, int, error) {
 	if rune(chunk[0]) != 'i' {
-		return 0, 0, false
+		return 0, 0, nil
 	}
 
 	num, err := strconv.Atoi(chunk[1:strings.IndexByte(chunk, 'e')])
 	if err != nil {
-		return 0, 0, false
+		return 0, 0, err
 	}
 
-	return num, len(strconv.Itoa(num)) + 2, true
+	return num, len(strconv.Itoa(num)) + 2, nil
 }
 
-func (d *Decoder) decodeSlice(chunk string) (any, int, bool) {
+func (d *Decoder) decodeSlice(chunk string) (any, int, error) {
 	if rune(chunk[0]) != 'l' {
-		return nil, 0, false
+		return nil, 0, nil
 	}
 
 	if rune(chunk[1]) == 'e' {
-		return []any{}, 2, true
+		return []any{}, 2, nil
 	}
 
-	slice := make([]any, 0, 1024)
+	result := make([]any, 0, 32)
 	offset := 1
 
 	for rune(chunk[offset]) != 'e' {
 		for _, f := range d.decoderFuncs {
-			if v, l, ok := f(chunk[offset:]); ok {
-				slice = append(slice, v)
+			if v, l, err := f(chunk[offset:]); l > 0 {
+				if err != nil {
+					return nil, 0, err
+				}
+				result = append(result, v)
 				offset += l
+				break
 			}
 		}
 	}
 
-	return slice, offset + 1, true
+	return result, offset + 1, nil
+}
+
+func (d *Decoder) decodeMap(chunk string) (any, int, error) {
+	if rune(chunk[0]) != 'd' {
+		return nil, 0, nil
+	}
+
+	if rune(chunk[1]) == 'e' {
+		return map[string]any{}, 2, nil
+	}
+
+	result := make(map[string]any, 32)
+	offset := 1
+
+	for rune(chunk[offset]) != 'e' {
+		key, keyLen, err := d.decodeString(chunk[offset:])
+		if err != nil {
+			return nil, 0, err
+		}
+
+		if keyLen == 0 {
+			return nil, 0, fmt.Errorf("invalid bencoded map key length")
+		}
+
+		for _, f := range d.decoderFuncs {
+			if v, l, err := f(chunk[offset+keyLen:]); l > 0 {
+				if err != nil {
+					return nil, 0, err
+				}
+
+				result[key.(string)] = v
+				offset += keyLen + l
+				break
+			}
+		}
+	}
+
+	return result, offset + 1, nil
 }
-- 
cgit v1.2.3