@@ -106,3 +106,148 @@ removeSquareBrackets <- function(code){
106106 }
107107 out
108108}
109+
110+ # Get for loop index range from a chunk of code
111+ forInfo <- function (x ){
112+ if (is.symbol(x )) return (NULL )
113+ if (x [[1 ]] != " for" return (NULL )
114+ as.list(x [[3 ]])
115+ }
116+
117+ # Attempt to combine for loops that share the same index range
118+ collapseLoopsInternal <- function (code ){
119+ if (is.symbol(code )) return (code )
120+ # Iterate over call components (except last one)
121+ for (i in 1 : (length(code )- 1 )){
122+ # Skip if just a symbol
123+ if (is.symbol(code [[i ]])) next
124+ # Skip if not a for loop
125+ if (code [[i ]][[1 ]] != " for" next
126+ # Get index range info from for loop
127+ i_info <- forInfo(code [[i ]])
128+ # Iterate over all subsequent call components after this one
129+ # looking for matching index ranges
130+ for (j in (i + 1 ): length(code )){
131+ # Skip if just a symbol
132+ if (is.symbol(code [[j ]])) next
133+ # Skip if not a for loop
134+ if (code [[j ]][[1 ]] != " for" next
135+ # Get index range info from the new for loop
136+ j_info <- forInfo(code [[j ]])
137+ # Check if the two index ranges match
138+ if (identical(i_info , j_info )){
139+ # Save existing for loop code for "parent" loop into new variable
140+ newloop <- code [[i ]]
141+ # Get loop index for this loop
142+ idx_i <- newloop [[2 ]]
143+ # Separate code inside loop
144+ internal <- newloop [[4 ]]
145+ # Drop the containing bracket from it
146+ internal [[1 ]] <- NULL
147+ # Get the code inside the "child" loop which will be added
148+ # to the parent loop
149+ # Note: could be a list if there is more than one line
150+ add_code <- code [[j ]][[4 ]]
151+ # Remove bracket
152+ add_code [[1 ]] <- NULL
153+ # Get loop index for the "child" loop
154+ idx_j <- code [[j ]][[2 ]]
155+ # Replace the existing index with the index from the parent loop
156+ add_code <- lapply(add_code , recursiveReplaceIndex , idx_j , idx_i )
157+ # Combine the parent and child loop code and insert it back into the for loop
158+ newloop [[4 ]] <- embedLinesInCurlyBrackets(c(internal , add_code ))
159+ # Insert the for loop into the full code
160+ code [[i ]] <- newloop
161+ # Mark the now-duplicated "child" code for removal later
162+ code [[j ]] <- " _REMOVE_"
163+ }
164+ }
165+ }
166+
167+ # Return all code parts except the stuff to be removed
168+ code [! sapply(code , function (x ) x == " _REMOVE_"
169+ }
170+
171+ # Recursively combine loops that share common index ranges
172+ collapseLoops <- function (code ){
173+ # Run the internal loop collapsing code once
174+ code <- collapseLoopsInternal(code )
175+ # Iterate over the result, looking for internal for loops and collapsing those
176+ if (is.call(code )){
177+ out <- lapply(code , function (x ){
178+ if (is.symbol(x )) return (x )
179+ if (x [[1 ]] == " for"
180+ x [[4 ]] <- collapseLoops(x [[4 ]])
181+ }
182+ x
183+ })
184+ out <- as.call(out )
185+ } else {
186+ out <- code
187+ }
188+ out
189+ }
190+
191+ # Replace complex indices (i_1, i_2, etc.) with a smaller number of
192+ # simplier indices (i, j, etc.) if possible
193+ simplifyIndices <- function (code , new_indices ){
194+ out <- lapply(code , function (x ){
195+ if (is.name(x ) | is.symbol(x )) return (x )
196+ if (x [[1 ]] == " for"
197+ unique_idx <- unique(extractAllIndices(x ))
198+ if (length(unique_idx ) > length(new_indices )){
199+ stop(" Not enough new indices provided" call. = FALSE )
200+ }
201+ for (i in 1 : length(unique_idx )){
202+ x <- replaceForLoopIndex(x , unique_idx [[i ]], new_indices [[i ]])
203+ }
204+ }
205+ x
206+ })
207+ as.call(out )
208+ }
209+
210+ # Replace the index in a loop recursively
211+ replaceForLoopIndex <- function (code , idx , newidx ){
212+ if (is.name(code ) | is.symbol(code )) return (code )
213+ if (code [[1 ]] == " for"
214+ if (code [[2 ]] == idx ) code [[2 ]] <- newidx
215+ code [[4 ]] <- recursiveReplaceIndex(code [[4 ]], idx , newidx )
216+ if (is.call(code [[4 ]])){
217+ code [[4 ]] <- as.call(lapply(code [[4 ]], function (x )
218+ replaceForLoopIndex(x , idx , newidx )))
219+ }
220+ }
221+ code
222+ }
223+
224+ # ' Simplify for loop structure in NIMBLE model code
225+ # '
226+ # ' Takes the code for a NIMBLE model and attempts to combine for loops
227+ # ' that share the same index range, in order to simplify the code
228+ # ' structure. Optionally, can also replace existing for loop indices
229+ # ' with a (potentially smaller, simpler) set of new indices.
230+ # ' This function is particularly useful for simplifying code generated by
231+ # ' macros, which often creates many for loops with the same indices and
232+ # ' uses complex indices like 'i_1', 'i_2', etc. which could be simplified to
233+ # ' 'i', 'j', etc.
234+ # '
235+ # ' @author Ken Kellner
236+ # '
237+ # ' @param code NIMBLE code for a model, such as from the output of model$getCode()
238+ # ' @param new_indices A list of new for loop indices that will replace the existing
239+ # ' indices. The new indices must be quoted values (i.e., "names"/symbols).
240+ # ' If NULL, letters starting with 'i' will be used. If FALSE, no indices will
241+ # ' be replaced.
242+ # '
243+ # ' @export
244+ simplifyForLoops <- function (code , new_indices = NULL ){
245+ out <- collapseLoops(code )
246+ if (is.null(new_indices )){
247+ new_indices <- lapply(letters [9 : 26 ], str2lang )
248+ } else if (is.logical(new_indices ) && ! new_indices ){
249+ return (out )
250+ }
251+ out <- simplifyIndices(out , new_indices )
252+ out
253+ }
0 commit comments