1/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
2
3Licensed under the Apache License, Version 2.0 (the "License");
4you may not use this file except in compliance with the License.
5You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9Unless required by applicable law or agreed to in writing, software
10distributed under the License is distributed on an "AS IS" BASIS,
11WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12See the License for the specific language governing permissions and
13limitations under the License.
14==============================================================================*/
15
16#include "tensorflow/lite/string_util.h"
17
18#include <stddef.h>
19#include <stdint.h>
20
21#include <cstdlib>
22#include <cstring>
23#include <vector>
24
25#include "tensorflow/lite/c/common.h"
26
27namespace tflite {
28
29void DynamicBuffer::AddString(const char* str, size_t len) {
30 data_.resize(data_.size() + len);
31 memcpy(data_.data() + offset_.back(), str, len);
32 offset_.push_back(offset_.back() + len);
33}
34
35void DynamicBuffer::AddString(const StringRef& string) {
36 AddString(string.str, string.len);
37}
38
39void DynamicBuffer::AddJoinedString(const std::vector<StringRef>& strings,
40 char separator) {
41 StringRef ref;
42 ref.str = &separator;
43 ref.len = 1;
44 AddJoinedString(strings, ref);
45}
46
47void DynamicBuffer::AddJoinedString(const std::vector<StringRef>& strings,
48 StringRef separator) {
49 // Resize the data buffer.
50 int total_len = (strings.size() - 1) * separator.len;
51 for (StringRef ref : strings) {
52 total_len += ref.len;
53 }
54 data_.resize(data_.size() + total_len);
55
56 char* dst = data_.data() + offset_.back();
57 for (size_t i = 0; i < strings.size(); ++i) {
58 // Fill separator if not first string.
59 if (i != 0) {
60 memcpy(dst, separator.str, separator.len);
61 dst += separator.len;
62 }
63
64 // Fill content of the string.
65 memcpy(dst, strings[i].str, strings[i].len);
66 dst += strings[i].len;
67 }
68 offset_.push_back(offset_.back() + total_len);
69}
70
71int DynamicBuffer::WriteToBuffer(char** buffer) {
72 // Allocate sufficient memory to tensor buffer.
73 int32_t num_strings = offset_.size() - 1;
74 // Total bytes include:
75 // * size of content (data_.size)
76 // * offset of each tensor (sizeof(int32_t) * num_strings)
77 // * length of whole buffer (int32_t)
78 // * num of strings (int32_t).
79 int32_t bytes = data_.size() // size of content
80 + sizeof(int32_t) * (num_strings + 2); // size of header
81
82 // Caller will take ownership of buffer.
83 *buffer = reinterpret_cast<char*>(malloc(bytes));
84
85 // Set num of string
86 memcpy(*buffer, &num_strings, sizeof(int32_t));
87
88 // Set offset of strings.
89 int32_t start = sizeof(int32_t) * (num_strings + 2);
90 for (size_t i = 0; i < offset_.size(); i++) {
91 int32_t offset = start + offset_[i];
92 memcpy(*buffer + sizeof(int32_t) * (i + 1), &offset, sizeof(int32_t));
93 }
94
95 // Copy data of strings.
96 memcpy(*buffer + start, data_.data(), data_.size());
97 return bytes;
98}
99
100#ifndef TF_LITE_STATIC_MEMORY
101void DynamicBuffer::WriteToTensorAsVector(TfLiteTensor* tensor) {
102 auto dims = TfLiteIntArrayCreate(1);
103 dims->data[0] = offset_.size() - 1; // Store number of strings.
104 WriteToTensor(tensor, dims);
105}
106
107void DynamicBuffer::WriteToTensor(TfLiteTensor* tensor,
108 TfLiteIntArray* new_shape) {
109 char* tensor_buffer;
110 int bytes = WriteToBuffer(&tensor_buffer);
111
112 if (new_shape == nullptr) {
113 new_shape = TfLiteIntArrayCopy(tensor->dims);
114 }
115
116 // Set tensor content pointer to tensor_buffer, and release original data.
117 TfLiteTensorReset(tensor->type, tensor->name, new_shape, tensor->params,
118 tensor_buffer, bytes, kTfLiteDynamic, tensor->allocation,
119 tensor->is_variable, tensor);
120}
121#endif // TF_LITE_STATIC_MEMORY
122
123int GetStringCount(const void* raw_buffer) {
124 // The first integers in the raw buffer is the number of strings.
125 return *static_cast<const int32_t*>(raw_buffer);
126}
127
128int GetStringCount(const TfLiteTensor* tensor) {
129 // The first integers in the raw buffer is the number of strings.
130 return GetStringCount(tensor->data.raw);
131}
132
133StringRef GetString(const void* raw_buffer, int string_index) {
134 const int32_t* offset =
135 static_cast<const int32_t*>(raw_buffer) + (string_index + 1);
136 return StringRef{
137 static_cast<const char*>(raw_buffer) + (*offset),
138 (*(offset + 1)) - (*offset),
139 };
140}
141
142StringRef GetString(const TfLiteTensor* tensor, int string_index) {
143 return GetString(tensor->data.raw, string_index);
144}
145
146} // namespace tflite
147