1// Copyright 2015 Jean Niklas L'orange.  All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package edn
6
7import (
8	"bytes"
9	"io"
10	"unicode/utf8"
11)
12
13var (
14	// we can't call it spaceBytes not to conflict with decode.go's spaceBytes.
15	spaceOutputBytes = []byte(" ")
16	commaOutputBytes = []byte(",")
17)
18
19func newline(dst io.Writer, prefix, indent string, depth int) {
20	dst.Write([]byte{'\n'})
21	dst.Write([]byte(prefix))
22	for i := 0; i < depth; i++ {
23		dst.Write([]byte(indent))
24	}
25}
26
27// Indent writes to dst an indented form of the EDN-encoded src. Each EDN
28// collection begins on a new, indented line beginning with prefix followed by
29// one or more copies of indent according to the indentation nesting. The data
30// written to dst does not begin with the prefix nor any indentation, and has
31// no trailing newline, to make it easier to embed inside other formatted EDN
32// data.
33//
34// Indent filters away whitespace, including comments and discards.
35func Indent(dst *bytes.Buffer, src []byte, prefix, indent string) error {
36	origLen := dst.Len()
37	err := IndentStream(dst, bytes.NewBuffer(src), prefix, indent)
38	if err != nil {
39		dst.Truncate(origLen)
40	}
41	return err
42}
43
44// IndentStream is an implementation of PPrint for generic readers and writers
45func IndentStream(dst io.Writer, src io.Reader, prefix, indent string) error {
46	var lex lexer
47	lex.reset()
48	tokStack := newTokenStack()
49	curType := tokenError
50	curSize := 0
51	d := NewDecoder(src)
52	depth := 0
53	for {
54		bs, tt, err := d.nextToken()
55		if err != nil {
56			return err
57		}
58		err = tokStack.push(tt)
59		if err != nil {
60			return err
61		}
62		prevType := curType
63		prevSize := curSize
64		if len(tokStack.toks) > 0 {
65			curType = tokStack.peek()
66			curSize = tokStack.peekCount()
67		}
68		switch tt {
69		case tokenMapStart, tokenVectorStart, tokenListStart, tokenSetStart:
70			if prevType == tokenMapStart {
71				dst.Write([]byte{' '})
72			} else if depth > 0 {
73				newline(dst, prefix, indent, depth)
74			}
75			dst.Write(bs)
76			depth++
77		case tokenVectorEnd, tokenListEnd, tokenMapEnd: // tokenSetEnd == tokenMapEnd
78			depth--
79			if prevSize > 0 { // suppress indent for empty collections
80				newline(dst, prefix, indent, depth)
81			}
82			// all of these are of length 1 in bytes, so utilise this for perf
83			dst.Write(bs)
84		case tokenTag:
85			// need to know what the previous type was.
86			switch prevType {
87			case tokenMapStart:
88				if prevSize%2 == 0 { // If previous size modulo 2 is equal to 0, we're a key
89					if prevSize > 0 {
90						dst.Write(commaOutputBytes)
91					}
92					newline(dst, prefix, indent, depth)
93				} else { // We're a value, add a space after the key
94					dst.Write(spaceOutputBytes)
95				}
96				dst.Write(bs)
97				dst.Write(spaceOutputBytes)
98			case tokenSetStart, tokenVectorStart, tokenListStart:
99				newline(dst, prefix, indent, depth)
100				dst.Write(bs)
101				dst.Write(spaceOutputBytes)
102			default: // tokenError or nested tag
103				dst.Write(bs)
104				dst.Write(spaceOutputBytes)
105			}
106		default:
107			switch prevType {
108			case tokenMapStart:
109				if prevSize%2 == 0 { // If previous size modulo 2 is equal to 0, we're a key
110					if prevSize > 0 {
111						dst.Write(commaOutputBytes)
112					}
113					newline(dst, prefix, indent, depth)
114				} else { // We're a value, add a space after the key
115					dst.Write(spaceOutputBytes)
116				}
117				dst.Write(bs)
118			case tokenSetStart, tokenVectorStart, tokenListStart:
119				newline(dst, prefix, indent, depth)
120				dst.Write(bs)
121			default: // toplevel or nested tag. This should collapse the whole tag tower
122				dst.Write(bs)
123			}
124		}
125		if tokStack.done() {
126			break
127		}
128	}
129	return nil
130}
131
132// PPrintOpts is a configuration map for PPrint. The values in this struct has
133// no effect as of now.
134type PPrintOpts struct {
135	RightMargin int
136	MiserWidth  int
137}
138
139func pprintIndent(dst io.Writer, shift int) {
140	spaces := make([]byte, shift+1)
141
142	spaces[0] = '\n'
143
144	// TODO: This may be slower than caching the size as a byte slice
145	for i := 1; i <= shift; i++ {
146		spaces[i] = ' '
147	}
148
149	dst.Write(spaces)
150}
151
152// PPrint writes to dst an indented form of the EDN-encoded src. This
153// implementation attempts to write idiomatic/readable EDN values, in a fashion
154// close to (but not quite equal to) clojure.pprint/pprint.
155//
156// PPrint filters away whitespace, including comments and discards.
157func PPrint(dst *bytes.Buffer, src []byte, opt *PPrintOpts) error {
158	origLen := dst.Len()
159	err := PPrintStream(dst, bytes.NewBuffer(src), opt)
160	if err != nil {
161		dst.Truncate(origLen)
162	}
163	return err
164}
165
166// PPrintStream is an implementation of PPrint for generic readers and writers
167func PPrintStream(dst io.Writer, src io.Reader, opt *PPrintOpts) error {
168	var lex lexer
169	var col, prevCollStart, curSize int
170	var prevColl bool
171
172	lex.reset()
173	tokStack := newTokenStack()
174
175	shift := make([]int, 1, 8) // pre-allocate some space
176	curType := tokenError
177	d := NewDecoder(src)
178
179	for {
180		bs, tt, err := d.nextToken()
181		if err != nil {
182			return err
183		}
184		err = tokStack.push(tt)
185		if err != nil {
186			return err
187		}
188		prevType := curType
189		prevSize := curSize
190		if len(tokStack.toks) > 0 {
191			curType = tokStack.peek()
192			curSize = tokStack.peekCount()
193		}
194		// Indentation
195		switch tt {
196		case tokenVectorEnd, tokenListEnd, tokenMapEnd:
197		default:
198			switch prevType {
199			case tokenMapStart:
200				if prevSize%2 == 0 && prevSize > 0 {
201					dst.Write(commaOutputBytes)
202					pprintIndent(dst, shift[len(shift)-1])
203					col = shift[len(shift)-1]
204				} else if prevSize%2 == 1 { // We're a value, add a space after the key
205					dst.Write(spaceOutputBytes)
206					col++
207				}
208			case tokenSetStart, tokenVectorStart, tokenListStart:
209				if prevColl {
210					// begin on new line where prevColl started
211					// This will look so strange for heterogenous maps.
212					pprintIndent(dst, prevCollStart)
213					col = prevCollStart
214				} else if prevSize > 0 {
215					dst.Write(spaceOutputBytes)
216					col++
217				}
218			}
219		}
220		switch tt {
221		case tokenMapStart, tokenVectorStart, tokenListStart, tokenSetStart:
222			dst.Write(bs)
223			col += len(bs)             // either 2 or 1
224			shift = append(shift, col) // we only use maps for now, but we'll utilise this more thoroughly later on
225		case tokenVectorEnd, tokenListEnd, tokenMapEnd: // tokenSetEnd == tokenMapEnd
226			dst.Write(bs) // all of these are of length 1 in bytes, so this is ok
227			prevCollStart = shift[len(shift)-1] - 1
228			shift = shift[:len(shift)-1]
229		case tokenTag:
230			bslen := utf8.RuneCount(bs)
231			dst.Write(bs)
232			dst.Write(spaceOutputBytes)
233			col += bslen + 1
234		default:
235			bslen := utf8.RuneCount(bs)
236			dst.Write(bs)
237			col += bslen
238		}
239		prevColl = (tt == tokenMapEnd || tt == tokenVectorEnd || tt == tokenListEnd)
240		if tokStack.done() {
241			break
242		}
243	}
244	return nil
245}
246