<commit_before>package jo

// NodeType indicates the type of the value stored in a Node.
type NodeType int

const (
	ObjectNode NodeType = iota
	ObjectKeyNode
	ArrayNode
	StringNode
	NumberNode
	BoolNode
	NullNode
)

// A Node is a part of a tree describing a JSON value. It contains
// information about the node's type, its original JSON representation
// (for literal values), and information about its position inside
// the tree.
type Node struct {
	Type NodeType
	Raw  []byte

	Parent                   *Node
	PrevSibling, NextSibling *Node
	FirstChild, LastChild    *Node
}

// AppendChild adds a child to the subject Node.
func (n *Node) AppendChild(c *Node) {
	if c.Parent != nil || c.PrevSibling != nil || c.NextSibling != nil {
		panic("jo: AppendChild called with an already attached Node")
	}

	if n.FirstChild == nil {
		n.FirstChild = c
	} else {
		n.LastChild.NextSibling = c
	}

	c.Parent = n
	c.PrevSibling = n.LastChild
	n.LastChild = c
}

// InsertChild adds a child to the subject Node, immediately before a
// reference node. Calling InsertChild with a nil reference node is
// functionally equivalent to just calling AppendChild.
func (n *Node) InsertChild(c, before *Node) {
	switch {
	case c.Parent != nil || c.PrevSibling != nil || c.NextSibling != nil:
		panic("jo: InsertChild called with an already attached Node")
	case before != nil && before.Parent != n:
		panic("jo: InsertChild called with unknown reference node")
	}

	c.Parent = n
	c.NextSibling = before

	if before == nil {
		if n.LastChild != nil {
			n.LastChild.NextSibling = c
		}
		c.PrevSibling = n.LastChild
		n.LastChild = c
	} else {
		if before.PrevSibling != nil {
			before.PrevSibling.NextSibling = c
		}
		c.PrevSibling = before.PrevSibling
		before.PrevSibling = c
	}

	if n.FirstChild == before {
		n.FirstChild = c
	}
}

// RemoveChild removes a child from the subject Node.
func (n *Node) RemoveChild(c *Node) {
	if c.Parent != n {
		panic("jo: RemoveChild called for non-child Node")
	}

	if n.FirstChild == c {
		n.FirstChild = c.NextSibling
	}
	if n.LastChild == c {
		n.LastChild = c.PrevSibling
	}

	if c.NextSibling != nil {
		c.NextSibling.PrevSibling = c.PrevSibling
	}
	if c.PrevSibling != nil {
		c.PrevSibling.NextSibling = c.NextSibling
	}

	c.Parent = nil
	c.PrevSibling = nil
	c.NextSibling = nil
}

// Parse parses raw JSON input and generates its tree representation.
func Parse(buf []byte) (*Node, error) {
	var p = parseState{buf: buf}

	root, err := p.close(p.scan())
	if err != nil {
		return nil, err
	}

	// consume all bytes after the top-level value
	if p.scan() == OpSyntaxError {
		return nil, p.s.LastError()
	}

	return root, nil
}

// parseState holds state when parsing a JSON value.
type parseState struct {
	s   Scanner
	buf []byte
	off int
}<commit_msg>Implement Node.MarshalJSON<commit_after>package jo

// NodeType indicates the type of the value stored in a Node.
type NodeType int

const (
	ObjectNode NodeType = iota
	ObjectKeyNode
	ArrayNode
	StringNode
	NumberNode
	BoolNode
	NullNode
)

// A Node is a part of a tree describing a JSON value. It contains
// information about the node's type, its original JSON representation
// (for literal values), and information about its position inside
// the tree.
type Node struct {
	Type NodeType
	Raw  []byte

	Parent                   *Node
	PrevSibling, NextSibling *Node
	FirstChild, LastChild    *Node
}

// AppendChild adds a child to the subject Node.
func (n *Node) AppendChild(c *Node) {
	if c.Parent != nil || c.PrevSibling != nil || c.NextSibling != nil {
		panic("jo: AppendChild called with an already attached Node")
	}

	if n.FirstChild == nil {
		n.FirstChild = c
	} else {
		n.LastChild.NextSibling = c
	}

	c.Parent = n
	c.PrevSibling = n.LastChild
	n.LastChild = c
}

// InsertChild adds a child to the subject Node, immediately before a
// reference node. Calling InsertChild with a nil reference node is
// functionally equivalent to just calling AppendChild.
func (n *Node) InsertChild(c, before *Node) {
	switch {
	case c.Parent != nil || c.PrevSibling != nil || c.NextSibling != nil:
		panic("jo: InsertChild called with an already attached Node")
	case before != nil && before.Parent != n:
		panic("jo: InsertChild called with unknown reference node")
	}

	c.Parent = n
	c.NextSibling = before

	if before == nil {
		if n.LastChild != nil {
			n.LastChild.NextSibling = c
		}
		c.PrevSibling = n.LastChild
		n.LastChild = c
	} else {
		if before.PrevSibling != nil {
			before.PrevSibling.NextSibling = c
		}
		c.PrevSibling = before.PrevSibling
		before.PrevSibling = c
	}

	if n.FirstChild == before {
		n.FirstChild = c
	}
}

// RemoveChild removes a child from the subject Node.
func (n *Node) RemoveChild(c *Node) {
	if c.Parent != n {
		panic("jo: RemoveChild called for non-child Node")
	}

	if n.FirstChild == c {
		n.FirstChild = c.NextSibling
	}
	if n.LastChild == c {
		n.LastChild = c.PrevSibling
	}

	if c.NextSibling != nil {
		c.NextSibling.PrevSibling = c.PrevSibling
	}
	if c.PrevSibling != nil {
		c.PrevSibling.NextSibling = c.NextSibling
	}

	c.Parent = nil
	c.PrevSibling = nil
	c.NextSibling = nil
}

// MarshalJSON marshals to JSON the tree structure rooted at the node.
// Note that there are no sanity checks, so malformed node trees will
// generate invalid JSON output.
func (n *Node) MarshalJSON() ([]byte, error) {
	out := make([]byte, n.marshalledSize())
	n.marshalTo(out)

	return out, nil
}

// marshalledSize returns the number of bytes required to marshal
// the node to JSON.
func (n *Node) marshalledSize() int {
	switch n.Type {
	case ArrayNode, ObjectNode:
		if n.FirstChild == nil {
			return 2
		}

		var total, count int
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			total += c.marshalledSize()
			count++
		}
		return 2 + total + (count - 1)

	case ObjectKeyNode:
		if c := n.FirstChild; c != nil {
			return len(n.Raw) + 1 + c.marshalledSize()
		}
		return len(n.Raw) + 1
	}

	return len(n.Raw)
}

// marshalTo writes the node's JSON representation to a byte slice,
// and returns the number of bytes written.
func (n *Node) marshalTo(dst []byte) (w int) {
	switch n.Type {
	case ObjectNode:
		dst[w] = '{'
		w++
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			if c != n.FirstChild {
				dst[w] = ','
				w++
			}
			w += c.marshalTo(dst[w:])
		}
		dst[w] = '}'
		w++

	case ArrayNode:
		dst[w] = '['
		w++
		for c := n.FirstChild; c != nil; c = c.NextSibling {
			if c != n.FirstChild {
				dst[w] = ','
				w++
			}
			w += c.marshalTo(dst[w:])
		}
		dst[w] = ']'
		w++

	case ObjectKeyNode:
		w += copy(dst, n.Raw)
		dst[w] = ':'
		w++
		if n.FirstChild != nil {
			w += n.FirstChild.marshalTo(dst[w:])
		}

	default:
		w = copy(dst[w:], n.Raw)
	}

	return w
}

// Parse parses raw JSON input and generates its tree representation.
func Parse(buf []byte) (*Node, error) {
	var p = parseState{buf: buf}

	root, err := p.close(p.scan())
	if err != nil {
		return nil, err
	}

	// consume all bytes after the top-level value
	if p.scan() == OpSyntaxError {
		return nil, p.s.LastError()
	}

	return root, nil
}

// parseState holds state when parsing a JSON value.
type parseState struct {
	s   Scanner
	buf []byte
	off int
}
{'fraction_non_alphanumeric': 0.11812881732370906, 'fraction_numerical': 0.000694058856191005, 'mean_word_length': 3.082152974504249, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '268369', 'n_tokens_mistral': 2658, 'n_tokens_neox': 2557, 'n_words': 824}