Fix mapPage method in PagedIterable for sync stack (Azure#35123)

* Fix mapPage method in PagedIterable for sync stack

* address pr comments

Author: srnagar

sdk/core/azure-core/src/main/java/com/azure/core/http/rest/PagedIterable.java

@@ -6,9 +6,11 @@
 import com.azure.core.util.IterableStream;
 import com.azure.core.util.paging.PageRetrieverSync;
 
+import java.util.ArrayList;
 import java.util.function.BiFunction;
 import java.util.function.Function;
 import java.util.function.Supplier;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 /**
@@ -73,6 +75,8 @@
  */
 public class PagedIterable<T> extends PagedIterableBase<T, PagedResponse<T>> {
     private final PagedFlux<T> pagedFlux;
+    private final Function<Integer, PagedResponse<T>> firstPageRetriever;
+    private final BiFunction<String, Integer, PagedResponse<T>> nextPageRetriever;
 
     /**
      * Creates instance given {@link PagedFlux}.
@@ -81,6 +85,8 @@ public class PagedIterable<T> extends PagedIterableBase<T, PagedResponse<T>> {
     public PagedIterable(PagedFlux<T> pagedFlux) {
         super(pagedFlux);
         this.pagedFlux = pagedFlux;
+        this.firstPageRetriever = null;
+        this.nextPageRetriever = null;
     }
 
     /**
@@ -150,7 +156,9 @@ public PagedIterable(Supplier<PagedResponse<T>> firstPageRetriever,
         this(() -> (continuationToken, pageSize) ->
             continuationToken == null
                  ? firstPageRetriever.get()
-                 : nextPageRetriever.apply(continuationToken), true);
+                 : nextPageRetriever.apply(continuationToken),
+            pageSize -> firstPageRetriever.get(),
+            (continuationToken, pageSize) -> nextPageRetriever.apply(continuationToken));
     }
 
     /**
@@ -175,7 +183,7 @@ public PagedIterable(Function<Integer, PagedResponse<T>> firstPageRetriever,
         BiFunction<String, Integer, PagedResponse<T>> nextPageRetriever) {
         this(() -> (continuationToken, pageSize) -> continuationToken == null
              ? firstPageRetriever.apply(pageSize)
-             : nextPageRetriever.apply(continuationToken, pageSize), true);
+             : nextPageRetriever.apply(continuationToken, pageSize), firstPageRetriever, nextPageRetriever);
     }
 
     /**
@@ -187,17 +195,51 @@ public PagedIterable(Function<Integer, PagedResponse<T>> firstPageRetriever,
      */
     @SuppressWarnings("deprecation")
     public <S> PagedIterable<S> mapPage(Function<T, S> mapper) {
-        return new PagedIterable<>(pagedFlux.mapPage(mapper));
+        if (pagedFlux != null) {
+            return new PagedIterable<>(pagedFlux.mapPage(mapper));
+        }
+        Function<Integer, PagedResponse<S>> firstMappedPageRetriever = pageSize -> {
+            PagedResponse<T> firstPageResponse = this.firstPageRetriever.apply(pageSize);
+            PagedResponse<S> firstMappedPageResponse = mapPagedResponse(firstPageResponse, mapper);
+            return firstMappedPageResponse;
+        };
+        BiFunction<String, Integer, PagedResponse<S>> nextMappedPageRetriever = (continuationToken, pageSize) -> {
+            PagedResponse<T> nextPageResponse = this.nextPageRetriever.apply(continuationToken, pageSize);
+            PagedResponse<S> nextMappedPageResponse = mapPagedResponse(nextPageResponse, mapper);
+            return nextMappedPageResponse;
+        };
+        return new PagedIterable<>(firstMappedPageRetriever, nextMappedPageRetriever);
     }
 
     /**
      * Create PagedIterable backed by Page Retriever Function Supplier.
      *
-     * @param provider the Page Retrieval Provider
-     * @param ignored param is ignored, exists in signature only to avoid conflict with first ctr
+     * @param provider the Page Retrieval Provider.
+     * @param firstPageRetriever first page retriever function to get the first page given the page size.
+     * @param nextPageRetriever next page retriever function to get the next page given a continuation token and the page size.
      */
-    private PagedIterable(Supplier<PageRetrieverSync<String, PagedResponse<T>>> provider, boolean ignored) {
+    private PagedIterable(Supplier<PageRetrieverSync<String, PagedResponse<T>>> provider,
+                          Function<Integer, PagedResponse<T>> firstPageRetriever,
+                          BiFunction<String, Integer, PagedResponse<T>> nextPageRetriever) {
         super(provider);
         this.pagedFlux = null;
+        this.firstPageRetriever = firstPageRetriever;
+        this.nextPageRetriever = nextPageRetriever;
     }
+
+    private <S> PagedResponse<S> mapPagedResponse(PagedResponse<T> pagedResponse, Function<T, S> mapper) {
+        if (pagedResponse == null) {
+            return null;
+        }
+        return new PagedResponseBase<String, S>(pagedResponse.getRequest(),
+            pagedResponse.getStatusCode(),
+            pagedResponse.getHeaders(),
+            pagedResponse.getValue()
+                .stream()
+                .map(mapper)
+                .collect(Collectors.toCollection(() -> new ArrayList<>(pagedResponse.getValue().size()))),
+            pagedResponse.getContinuationToken(),
+            null);
+    }
+
 }

sdk/core/azure-core/src/test/java/com/azure/core/http/rest/PagedIterableTest.java

@@ -148,6 +148,19 @@ public void streamByPageMap(int numberOfPages) {
         }
     }
 
+    @ParameterizedTest
+    @ValueSource(ints = {0, 1, 2, 3, 4, 5})
+    public void streamByPageMapSync(int numberOfPages) {
+        PagedIterable<Integer> integerPagedIterable = getIntegerPagedIterable(numberOfPages);
+        List<PagedResponse<String>> pages = integerPagedIterable.mapPage(String::valueOf).streamByPage()
+            .collect(Collectors.toList());
+
+        assertEquals(numberOfPages, pages.size());
+        for (int i = 0; i < numberOfPages; i++) {
+            assertEquals(pagedStringResponses.get(i).getValue(), pages.get(i).getValue());
+        }
+    }
+
     @ParameterizedTest
     @ValueSource(ints = {0, 5})
     public void iterateByPageMap(int numberOfPages) {
@@ -162,6 +175,20 @@ public void iterateByPageMap(int numberOfPages) {
         }
     }
 
+    @ParameterizedTest
+    @ValueSource(ints = {0, 1, 2, 3, 4, 5})
+    public void iterateByPageMapSync(int numberOfPages) {
+        PagedIterable<Integer> integerPagedIterable = getIntegerPagedIterable(numberOfPages);
+
+        List<PagedResponse<String>> pages = new ArrayList<>();
+        integerPagedIterable.mapPage(String::valueOf).iterableByPage().iterator().forEachRemaining(pages::add);
+
+        assertEquals(numberOfPages, pages.size());
+        for (int i = 0; i < numberOfPages; i++) {
+            assertEquals(pagedStringResponses.get(i).getValue(), pages.get(i).getValue());
+        }
+    }
+
     @ParameterizedTest
     @ValueSource(ints = {0, 5})
     public void streamByTMap(int numberOfPages) {
@@ -174,6 +201,17 @@ public void streamByTMap(int numberOfPages) {
             .collect(Collectors.toList()), values);
     }
 
+    @ParameterizedTest
+    @ValueSource(ints = {0, 1, 2, 3, 4, 5})
+    public void streamByTMapSync(int numberOfPages) {
+        PagedIterable<Integer> pagedIterable = getIntegerPagedIterable(numberOfPages);
+        List<String> values = pagedIterable.mapPage(String::valueOf).stream().collect(Collectors.toList());
+
+        assertEquals(numberOfPages * 3, values.size());
+        assertEquals(Stream.iterate(0, i -> i + 1).limit(numberOfPages * 3L).map(String::valueOf)
+            .collect(Collectors.toList()), values);
+    }
+
     @ParameterizedTest
     @ValueSource(ints = {0, 5})
     public void iterateByTMap(int numberOfPages) {
@@ -187,6 +225,18 @@ public void iterateByTMap(int numberOfPages) {
             .collect(Collectors.toList()), values);
     }
 
+    @ParameterizedTest
+    @ValueSource(ints = {0, 5})
+    public void iterateByTMapSync(int numberOfPages) {
+        PagedIterable<Integer> pagedIterable = getIntegerPagedIterable(numberOfPages);
+        List<String> values = new ArrayList<>();
+        pagedIterable.mapPage(String::valueOf).iterator().forEachRemaining(values::add);
+
+        assertEquals(numberOfPages * 3, values.size());
+        assertEquals(Stream.iterate(0, i -> i + 1).limit(numberOfPages * 3L).map(String::valueOf)
+            .collect(Collectors.toList()), values);
+    }
+
     @Test
     public void streamFirstPage() {
         TestPagedFlux<Integer> pagedFlux = getTestPagedFlux();