Add "buffer management" and "shrink" API calls to vstr.

vstr is initially intended to deal with arbitrary-length strings. By
providing a bit lower-level API calls, it will be also useful to deal
with arbitrary-length I/O buffers (the difference from strings is that
buffers are filled from "outside", via I/O).

Another issue, especially aggravated by I/O buffer use, is alloc size
vs actual size length. If allocated 1Mb for buffer, but actually
read 1 byte, we don't want to keep rest of 1Mb be locked by this I/O
result, but rather return it to heap ASAP ("shrink" buffer before passing
it to qstr_from_str_take()).

3 files changed, 48 insertions, 5 deletions

diff --git a/py/lexer.c b/py/lexer.c
index f7f9c631f3..da8967b163 100644
--- a/py/lexer.c
+++ b/py/lexer.c
@@ -614,7 +614,7 @@ mp_lexer_t *mp_lexer_new(const char *src_name, void *stream_data, mp_lexer_strea
     lex->num_indent_level = 1;
     lex->indent_level = m_new(uint16_t, lex->alloc_indent_level);
     lex->indent_level[0] = 0;
-    vstr_init(&lex->vstr);
+    vstr_init(&lex->vstr, 32);
 
     // preload characters
     lex->chr0 = stream_next_char(stream_data);
diff --git a/py/misc.h b/py/misc.h
index 1bf4d8f291..2d6fb978ac 100644
--- a/py/misc.h
+++ b/py/misc.h
@@ -58,15 +58,19 @@ typedef struct _vstr_t {
     bool had_error;
 } vstr_t;
 
-void vstr_init(vstr_t *vstr);
+void vstr_init(vstr_t *vstr, int alloc);
 void vstr_clear(vstr_t *vstr);
 vstr_t *vstr_new(void);
+vstr_t *vstr_new_size(int alloc);
 void vstr_free(vstr_t *vstr);
 void vstr_reset(vstr_t *vstr);
 bool vstr_had_error(vstr_t *vstr);
 char *vstr_str(vstr_t *vstr);
 int vstr_len(vstr_t *vstr);
 void vstr_hint_size(vstr_t *vstr, int size);
+char  *vstr_extend(vstr_t *vstr, int size);
+bool vstr_set_size(vstr_t *vstr, int size);
+bool vstr_shrink(vstr_t *vstr);
 char *vstr_add_len(vstr_t *vstr, int len);
 void vstr_add_byte(vstr_t *vstr, byte v);
 void vstr_add_char(vstr_t *vstr, unichar chr);
diff --git a/py/vstr.c b/py/vstr.c
index 80841b24ca..18e6d2d0a8 100644
--- a/py/vstr.c
+++ b/py/vstr.c
@@ -6,8 +6,8 @@
 // returned value is always at least 1 greater than argument
 #define ROUND_ALLOC(a) (((a) & ((~0) - 7)) + 8)
 
-void vstr_init(vstr_t *vstr) {
-    vstr->alloc = 32;
+void vstr_init(vstr_t *vstr, int alloc) {
+    vstr->alloc = alloc;
     vstr->len = 0;
     vstr->buf = m_new(char, vstr->alloc);
     if (vstr->buf == NULL) {
@@ -28,7 +28,16 @@ vstr_t *vstr_new(void) {
     if (vstr == NULL) {
         return NULL;
     }
-    vstr_init(vstr);
+    vstr_init(vstr, 32);
+    return vstr;
+}
+
+vstr_t *vstr_new_size(int alloc) {
+    vstr_t *vstr = m_new(vstr_t, 1);
+    if (vstr == NULL) {
+        return NULL;
+    }
+    vstr_init(vstr, alloc);
     return vstr;
 }
 
@@ -63,6 +72,36 @@ int vstr_len(vstr_t *vstr) {
     return vstr->len;
 }
 
+// Extend vstr strictly to by requested size, return pointer to newly added chunk
+char  *vstr_extend(vstr_t *vstr, int size) {
+    char *new_buf = m_renew(char, vstr->buf, vstr->alloc, vstr->alloc + size);
+    if (new_buf == NULL) {
+        vstr->had_error = true;
+        return NULL;
+    }
+    char *p = new_buf + vstr->alloc;
+    vstr->alloc += size;
+    vstr->buf = new_buf;
+    return p;
+}
+
+// Shrink vstr to be given size
+bool vstr_set_size(vstr_t *vstr, int size) {
+    char *new_buf = m_renew(char, vstr->buf, vstr->alloc, size);
+    if (new_buf == NULL) {
+        vstr->had_error = true;
+        return false;
+    }
+    vstr->buf = new_buf;
+    vstr->alloc = vstr->len = size;
+    return true;
+}
+
+// Shrink vstr allocation to its actual length
+bool vstr_shrink(vstr_t *vstr) {
+    return vstr_set_size(vstr, vstr->len);
+}
+
 bool vstr_ensure_extra(vstr_t *vstr, int size) {
     if (vstr->len + size + 1 > vstr->alloc) {
         int new_alloc = ROUND_ALLOC((vstr->len + size + 1) * 2);