cleanup: use helper methods from secretness analysis

lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.cpp

@@ -224,6 +224,12 @@ bool isSecret(Value value, DataFlowSolver *solver) {
   return isSecret(lattice);
 }
 
+bool isInitialized(Value value, DataFlowSolver *solver) {
+  auto *lattice = solver->lookupState<SecretnessLattice>(value);
+  if (!lattice) return false;
+  return lattice->getValue().isInitialized();
+}
+
 bool isSecret(const SecretnessLattice *lattice) {
   if (!lattice) {
     return false;

lib/Analysis/SecretnessAnalysis/SecretnessAnalysis.h

@@ -226,6 +226,8 @@ void annotateSecretness(Operation *top, DataFlowSolver *solver, bool verbose);
 // analysis
 bool isSecret(Value value, DataFlowSolver *solver);
 
+bool isInitialized(Value value, DataFlowSolver *solver);
+
 bool isSecret(const SecretnessLattice *lattice);
 
 bool isSecret(ValueRange values, DataFlowSolver *solver);

lib/Transforms/ConvertIfToSelect/ConvertIfToSelect.cpp

@@ -38,12 +38,10 @@ struct IfToSelectConversion : OpRewritePattern<scf::IfOp> {
 
   LogicalResult matchAndRewrite(scf::IfOp ifOp,
                                 PatternRewriter &rewriter) const override {
-    auto *lattice = solver->lookupState<SecretnessLattice>(ifOp.getOperand());
-    Secretness secretness = lattice ? lattice->getValue() : Secretness();
-
     // Convert ops with "secret" and, conservatively, "unknown" (uninitialized)
     // conditions but skip conversion if the condition is known to be non-secret
-    if (secretness.isInitialized() && !secretness.getSecretness())
+    auto ifOperand = ifOp.getOperand();
+    if (isInitialized(ifOperand, solver) && !isSecret(ifOperand, solver))
       return failure();
 
     // Hoist instructions in the 'then' and 'else' regions

lib/Transforms/ConvertSecretExtractToStaticExtract/ConvertSecretExtractToStaticExtract.cpp

@@ -49,25 +49,17 @@ struct SecretExtractToStaticExtractConversion
 
     auto index = extractOp.getIndices().front();
 
-    auto *indexSecretnessLattice =
-        solver->lookupState<SecretnessLattice>(index);
-
-    // Use secretness from lattice or, if no lattice found, set to unknown
-    auto indexSecretness = indexSecretnessLattice
-                               ? indexSecretnessLattice->getValue()
-                               : Secretness();
-
     // If lattice is set to unknown,
     // apply transformation anyway but emit a warning
-    if (!indexSecretness.isInitialized()) {
+    bool initialized = isInitialized(index, solver);
+    if (!initialized) {
       extractOp->emitWarning()
           << "Secretness for tensor.extract index is unknown. "
              "Conservatively, the transformation will be applied:";
     }
 
     // If index is known to be public, no transformation is needed
-    if (indexSecretness.isInitialized() && !indexSecretness.getSecretness())
-      return failure();
+    if (initialized && !isSecret(index, solver)) return failure();
 
     ImplicitLocOpBuilder builder(extractOp->getLoc(), rewriter);
 

lib/Transforms/ConvertSecretInsertToStaticInsert/ConvertSecretInsertToStaticInsert.cpp

@@ -50,25 +50,19 @@ struct SecretInsertToStaticInsertConversion
     auto tensor = insertOp.getDest();
     auto insertedValue = insertOp.getScalar();
 
-    auto *indexSecretnessLattice =
-        solver->lookupState<SecretnessLattice>(index);
-
-    // use secretness from lattice or, if no lattice found, set to unknown
-    auto indexSecretness = indexSecretnessLattice
-                               ? indexSecretnessLattice->getValue()
-                               : Secretness();
-
-    // If lattice is set to unknown,
-    // apply transformation anyway but emit a warning
-    if (!indexSecretness.isInitialized()) {
+    // If lattice is set to unknown, apply transformation anyway but emit a
+    // warning
+    bool initialized = isInitialized(index, solver);
+    if (!initialized) {
       insertOp->emitWarning()
           << "Secretness for tensor.insert index is unknown. "
              "Conservatively, the transformation will be applied:";
     }
 
     // If index is known to be public, no transformation is needed
-    if (indexSecretness.isInitialized() && !indexSecretness.getSecretness())
+    if (initialized && !isSecret(index, solver)) {
       return failure();
+    }
 
     ImplicitLocOpBuilder builder(insertOp->getLoc(), rewriter);
 

lib/Transforms/ConvertSecretWhileToStaticFor/ConvertSecretWhileToStaticFor.cpp

@@ -63,21 +63,15 @@ struct SecretWhileToStaticForConversion : OpRewritePattern<scf::WhileOp> {
       counter++;
     }
 
-    auto *conditionSecretnessLattice =
-        solver->lookupState<SecretnessLattice>(whileCondition);
-
-    if (!conditionSecretnessLattice) {
+    if (!isInitialized(whileCondition, solver)) {
       InFlightDiagnostic diag =
           whileOp.emitWarning()
           << "Secretness for scf.while condition has not been set";
       return failure();
     }
 
-    bool isConditionSecret =
-        conditionSecretnessLattice->getValue().getSecretness();
-
     // If condition is not secret, no transformation is needed
-    if (!isConditionSecret) {
+    if (!isSecret(whileCondition, solver)) {
       return failure();
     }