From 15965b0b054bcd58182f08cd74401a7433a420b4 Mon Sep 17 00:00:00 2001
From: Ozzie Gooen <ozzieagooen@gmail.com>
Date: Sun, 15 May 2022 17:24:54 -0400
Subject: [PATCH] First simple version of samplesMap

---
 .../Distributions/SampleSetDist/SampleSetDist.res | 11 +++++++++++
 .../Reducer_Dispatch/Reducer_Dispatch_BuiltIn.res | 15 +++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/packages/squiggle-lang/src/rescript/Distributions/SampleSetDist/SampleSetDist.res b/packages/squiggle-lang/src/rescript/Distributions/SampleSetDist/SampleSetDist.res
index 55b334c5..5426fd5e 100644
--- a/packages/squiggle-lang/src/rescript/Distributions/SampleSetDist/SampleSetDist.res
+++ b/packages/squiggle-lang/src/rescript/Distributions/SampleSetDist/SampleSetDist.res
@@ -83,6 +83,17 @@ let sampleN = (t: t, n) => {
   }
 }
 
+let samplesMap = (~fn: (float) => result<float, Operation.Error.t>, t: t): result<
+  t,
+  Operation.Error.t,
+> => {
+  let samples = T.get(t)->E.A2.fmap(fn)
+
+  E.A.R.firstErrorOrOpen(samples)->E.R2.fmap(x =>
+    E.R.toExn("Input of samples should be larger than 5", make(x))
+  )
+}
+
 //TODO: Figure out what to do if distributions are different lengths. ``zip`` is kind of inelegant for this.
 let map2 = (~fn: (float, float) => result<float, Operation.Error.t>, ~t1: t, ~t2: t): result<
   t,
diff --git a/packages/squiggle-lang/src/rescript/Reducer/Reducer_Dispatch/Reducer_Dispatch_BuiltIn.res b/packages/squiggle-lang/src/rescript/Reducer/Reducer_Dispatch/Reducer_Dispatch_BuiltIn.res
index d1a77040..b1deac6d 100644
--- a/packages/squiggle-lang/src/rescript/Reducer/Reducer_Dispatch/Reducer_Dispatch_BuiltIn.res
+++ b/packages/squiggle-lang/src/rescript/Reducer/Reducer_Dispatch/Reducer_Dispatch_BuiltIn.res
@@ -99,6 +99,19 @@ let callInternal = (call: functionCall, environment, reducer: ExpressionT.reduce
     rMappedList->Result.map(mappedList => mappedList->Belt.List.toArray->EvArray)
   }
 
+  let doMapSampleSetDist = (sampleSetDist: SampleSetDist.t, aLambdaValue) => {
+    let fn = r =>
+      switch Lambda.doLambdaCall(aLambdaValue, list{EvNumber(r)}, environment, reducer) {
+      | Ok(EvNumber(f)) => Ok(f)
+      | _ => Error(Operation.NotYetImplemented)
+      }
+    let newDist = SampleSetDist.samplesMap(~fn, sampleSetDist)
+    switch newDist {
+    | Ok(r) => Ok(EvDistribution(SampleSet(r)))
+    | Error(r) => Error(RETodo(""))
+    }
+  }
+
   let doReduceArray = (aValueArray, initialValue, aLambdaValue) => {
     aValueArray->Belt.Array.reduce(Ok(initialValue), (rAcc, elem) =>
       rAcc->Result.flatMap(acc =>
@@ -128,6 +141,8 @@ let callInternal = (call: functionCall, environment, reducer: ExpressionT.reduce
   | ("keep", [EvArray(aValueArray), EvLambda(aLambdaValue)]) =>
     doKeepArray(aValueArray, aLambdaValue)
   | ("map", [EvArray(aValueArray), EvLambda(aLambdaValue)]) => doMapArray(aValueArray, aLambdaValue)
+  | ("mapSamples", [EvDistribution(SampleSet(dist)), EvLambda(aLambdaValue)]) =>
+    doMapSampleSetDist(dist, aLambdaValue)
   | ("reduce", [EvArray(aValueArray), initialValue, EvLambda(aLambdaValue)]) =>
     doReduceArray(aValueArray, initialValue, aLambdaValue)
   | ("reduceReverse", [EvArray(aValueArray), initialValue, EvLambda(aLambdaValue)]) =>