diff --git a/tree/treeplayer/src/TTreeFormula.cxx b/tree/treeplayer/src/TTreeFormula.cxx index fe3c7366ec7ec..a883456647be6 100644 --- a/tree/treeplayer/src/TTreeFormula.cxx +++ b/tree/treeplayer/src/TTreeFormula.cxx @@ -5526,6 +5526,40 @@ void TTreeFormula::ResetDimensions() { //} } + + // Handle the case of an Alt$(primary,alternate) expression that provides + // the only source of multiplicity, i.e. it's used not alongside another + // array that drives the iteration. Example: Scan("Alt$(x,-1)"), instead of + // Scan("y-Alt$(x,-1)"). In that case the number of instances must be + // dictated by the `primary` array, so we register it with the manager. We + // only get here when fMultiplicity == 0, i.e. nothing else in this formula + // drives an iteration. + if (fMultiplicity == 0) { + for (i = 0; i < fNoper; i++) { + Int_t action = GetAction(i); + if (action == kAlternate || action == kAlternateString) { + TTreeFormula *primary = static_cast(fAliases.UncheckedAt(i)); + R__ASSERT(primary); + switch (primary->GetManager()->GetMultiplicity()) { + case 1: // primary is a variable length array: drive a variable length loop + fMultiplicity = 1; + fManager->Add(primary); + break; + case 2: // primary is a fixed length array: drive a fixed length loop + if (fMultiplicity != 1) + fMultiplicity = 2; + fManager->Add(primary); + break; + // multiplicity 0 (scalar) or -1 (0-or-1 element): nothing to loop + // over, leave fMultiplicity at 0 + } + // An Alt$(primary,alternate) expression occupies two consecutive + // operands: the `primary` action we just handled, immediately + // followed by the `alternate`, which we can skip. + ++i; + } + } + } } //////////////////////////////////////////////////////////////////////////////// diff --git a/tree/treeplayer/test/regressions.cxx b/tree/treeplayer/test/regressions.cxx index c174a9c5c2039..ab08da89f15fd 100644 --- a/tree/treeplayer/test/regressions.cxx +++ b/tree/treeplayer/test/regressions.cxx @@ -507,6 +507,89 @@ TEST(TTreeScan, chainNameWithDifferentTreeName) } } +// Alt$(primary, alternate) used on its own, with 'primary' a variable-length +// array, must loop over all the elements of the array instead of just the +// first one. Previously the array dimension of 'primary' was not propagated to +// the enclosing TTreeFormula manager when Alt$ was the only source of +// multiplicity, so the formula reported a single instance per entry. +TEST(TTreeFormulaRegressions, AltDollarVariableArray) +{ + TTree t("t", "t"); + Int_t n = 3; + Float_t x[3]{}; + t.Branch("n", &n); + t.Branch("x", &x, "x[n]/F"); + x[0] = 1; + x[1] = 2; + x[2] = 3; + t.Fill(); + x[0] = 4; + x[1] = 5; + x[2] = 6; + t.Fill(); + n = 2; + x[0] = -1; + x[1] = -2; + x[2] = -3; + t.Fill(); + n = 1; + x[0] = 0; + t.Fill(); + + auto evalAll = [](TTree &tree, TTreeFormula &form, Long64_t entry) { + tree.LoadTree(entry); + const Int_t ndata = form.GetNdata(); + std::vector values; + for (Int_t i = 0; i < ndata; ++i) + values.push_back(form.EvalInstance(i)); + return values; + }; + + // 1) Standalone Alt$ over the whole array loops over its elements. + { + TTreeFormula form("form", "Alt$(x,-1)", &t); + EXPECT_EQ(form.GetMultiplicity(), 1); + const std::vector> expected{{1, 2, 3}, {4, 5, 6}, {-1, -2}, {0}}; + for (Long64_t entry = 0; entry < t.GetEntries(); ++entry) + EXPECT_EQ(evalAll(t, form, entry), expected[entry]) << "entry " << entry; + } + + // 2) A fixed index into the variable-length array stays a single value per + // entry and falls back to the alternate when the index is out of range. + // This is the documented Alt$(arr[i], default) use case and must not be + // turned into a zero-instance (dropped) entry. + { + TTreeFormula form("form", "Alt$(x[2],-1)", &t); + const std::vector> expected{{3}, {6}, {-1}, {-1}}; // x[2] only exists for n==3 + for (Long64_t entry = 0; entry < t.GetEntries(); ++entry) + EXPECT_EQ(evalAll(t, form, entry), expected[entry]) << "entry " << entry; + } +} + +// Companion to AltDollarVariableArray: when Alt$ is combined with another array +// that drives the iteration, the alternate must pad the shorter array rather +// than shrinking the loop. See the TTree::Draw documentation for Alt$. +TEST(TTreeFormulaRegressions, AltDollarPadsShorterArray) +{ + TTree t("t", "t"); + Int_t n1 = 3, n2 = 2; + Float_t a1[3]{10, 20, 30}; + Float_t a2[3]{1, 2, 3}; + t.Branch("n1", &n1); + t.Branch("n2", &n2); + t.Branch("a1", &a1, "a1[n1]/F"); + t.Branch("a2", &a2, "a2[n2]/F"); + t.Fill(); + + // The loop is driven by a1 (3 elements); Alt$(a2,0) yields a2[0],a2[1],0. + TTreeFormula form("form", "a1+Alt$(a2,0)", &t); + t.LoadTree(0); + ASSERT_EQ(form.GetNdata(), 3); + const std::vector expected{11, 22, 30}; + for (Int_t i = 0; i < 3; ++i) + EXPECT_FLOAT_EQ(form.EvalInstance(i), expected[i]) << "instance " << i; +} + // https://github.com/root-project/root/issues/20249 TEST(TTreeScan, TTreeGetBranchOfFriendTChain) {